Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
package main
import (
"flag"
"fmt"
"net"
"os"
"sync"
"time"
)
var addr string
var closeChan = make(chan struct{}, 1024)
var messageChan = make(chan net.Conn, 1024)
var feedbackChan = make(chan net.Conn, 1024)
func main() {
flag.StringVar(&addr, "addr", ":1202", "listen addr")
flag.Parse()
l, err := net.Listen("tcp", addr)
if err != nil {
os.Exit(1)
}
defer l.Close()
connMap := make(map[net.Conn]bool)
mutex := new(sync.RWMutex)
go func() {
for {
select {
case conn := <-feedbackChan:
mutex.Lock()
delete(connMap, conn)
mutex.Unlock()
default:
mutex.RLock()
for conn := range connMap {
messageChan <- conn
}
mutex.RUnlock()
}
}
}()
go func() {
for {
conn, err := l.Accept()
if err != nil {
os.Exit(1)
}
mutex.Lock()
connMap[conn] = true
mutex.Unlock()
}
}()
for {
var i int
_, err := fmt.Scanf("%d", &i)
if err != nil {
fmt.Println("please input an integer", err.Error())
continue
}
if i <= 0 {
shrinkGoRoutines(-i)
} else {
expandGoRoutines(i)
}
}
}
func expandGoRoutines(num int) {
var i int
for i = 0; i < num; i++ {
go func() {
buf := make([]byte, 1024)
for {
select {
case conn := <-messageChan:
conn.SetReadDeadline(time.Now().Add(time.Microsecond * 50))
len, err := conn.Read(buf)
E, ok := err.(net.Error)
if ok && E.Timeout() {
continue
}
if err != nil {
conn.Close()
feedbackChan <- conn
continue
}
conn.Write(append([]byte(addr+" "), buf[:len]...))
case <-closeChan:
return
}
}
}()
}
}
func shrinkGoRoutines(num int) {
var i int
for i = 0; i < num; i++ {
closeChan <- struct{}{}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment