一个简单的go的Socket服务器工作池

一个简单的go的Socket服务器工作池

处理收包和发包是工作池

Server.go

package main

import (
   "fmt"
   "net"
   "strconv"
   "time"
   "sync"
)

type Task struct {
   conn *Session
   data []byte
}

type Result struct {
   conn *Session
   data []byte
}

var tasks = make(chan Task, 10)
var results = make(chan Result, 10)

type Session struct {
   conn net.Conn
   lock sync.Mutex
}

//处理接收的客户端包，并返回返回包
func process(data []byte) []byte {
   fmt.Println(string(data))
   return []byte("return from server.")
}

//接收发送包工作池
func worker(wg *sync.WaitGroup) {
   defer wg.Done()

   for {
      select {
      case task := <- tasks://接收的包
         fmt.Println("work get task...%s", task.data)
         task.conn.lock.Lock()//加锁为了保证同一个客户端包是顺序处理的
         result := Result{task.conn, process(task.data)}
         results <- result
         task.conn.lock.Unlock()
      case result := <- results://要发送的包
         fmt.Printf("Task id %s \n", string(result.data))
         result.conn.lock.Lock()//加锁为了保证同一个客户端包是顺序处理的
         result.conn.conn.Write(result.data)
         result.conn.lock.Unlock()
      }
   }
}

//创建工作池
func createWorkerPool(numOfWorkers int) {
   var wg sync.WaitGroup

   for i:= 0; i < numOfWorkers; i++ {
      wg.Add(1)
      go worker(&wg)
   }

   wg.Wait()
   close(results)
}

//将任务放到工作池
func allocate(data []byte, conn *Session) {
   task := Task{conn, data}
   tasks <- task

   time.Sleep(1*time.Second)
}

//启动工作池
func startPool() {
   startTime := time.Now()

   numOfWorkers := 20

   go createWorkerPool(numOfWorkers)

   endTime := time.Now()
   diff := endTime.Sub(startTime)

   fmt.Printf("total time taken :", diff.Seconds(), " seconds")
}

//启动监听
func StartListen(port int) {
   sport := ":"
   sport += strconv.Itoa(port)
   fmt.Println("listen port :", sport)
   listener, err :=  net.Listen("tcp", sport)
   if(err != nil) {
      fmt.Println("listen err:", err)
      return
   }

   for{
      fmt.Println("begin to accept")
      conn, err := listener.Accept()
      if err == nil {
         fmt.Println("accept a client")
         sess := new(Session)
         sess.conn = conn
         go HandleConn(sess)
      }
   }

   defer listener.Close()
}

/**
 * 接收数据包
 */
func RecvDatan(conn net.Conn, len int) ([]byte, bool) {
   //阻塞
   conn.SetReadDeadline(time.Now().Add(3*time.Second))

   var nTotalRecv int
   buff := make([]byte, len)
   nTotalRecv = 0
   for {
      if nTotalRecv >= len {
         break
      }
      nRecv, err := conn.Read(buff[nTotalRecv:])
      if err != nil {
         fmt.Println("connect read conn failed.", err, " nRecv:", nRecv, ",totalrecv:", nTotalRecv)
         return buff, false
      }
      nTotalRecv += nRecv
   }
   return buff, true
}

//收包协程
func HandleConn(conn *Session) {
   for{
      buff, bRet := RecvDatan(conn.conn, 8)
      if bRet == true {
         fmt.Println("client:", conn.conn.RemoteAddr(), string(buff))
         allocate(buff, conn)
      }
   }
}

func main() {
   startPool()
   StartListen(2020)

}