gnet-play

server.go

package server

import (
	"context"
	"log"
	"time"

	"github.com/panjf2000/gnet/v2"
)

// Config is the configuration for the WSServer.
type Config struct {
	NumLoops     int           `mapstructure:"num_loops" json:"num_loops" default:"0"`
	MultiCore    bool          `mapstructure:"multi_core" json:"multi_core" default:"false"`
	KeepAlive    time.Duration `mapstructure:"keep_alive" json:"keep_alive" default:"1h"`
	LockThread   bool          `mapstructure:"lock_thread" json:"lock_thread" default:"false"`
	ReadBufCap   int           `mapstructure:"read_buf_cap" json:"read_buf_cap" default:"4098"`
	WriteBufCap  int           `mapstructure:"write_buf_cap" json:"write_buf_cap" default:"65536"`
	SockRecvBuf  int           `mapstructure:"sock_recv_buf" json:"sock_recv_buf" default:"4096"`
	SockSendBuf  int           `mapstructure:"sock_send_buf" json:"sock_send_buf" default:"4096"`
	AntsPoolSize int           `mapstructure:"ants_pool_size" json:"ants_pool_size" default:"10000"`
	AsyncWrite   bool          `mapstructure:"async_write" json:"async_write" default:"true"`
}

// Serve starts the websocket connection manager.
func Serve(ctx context.Context, addr string, cfg Config) error {
	echo := &SockServer{
		ctx:  ctx,
		cfg:  cfg,
	}

	return gnet.Run(echo, addr,
		gnet.WithNumEventLoop(cfg.NumLoops),
		gnet.WithTicker(true),
		gnet.WithReadBufferCap(cfg.ReadBufCap),
		gnet.WithWriteBufferCap(cfg.WriteBufCap),
		gnet.WithTCPKeepAlive(cfg.KeepAlive),
		gnet.WithSocketSendBuffer(cfg.SockSendBuf),
		gnet.WithSocketRecvBuffer(cfg.SockRecvBuf),
		gnet.WithLockOSThread(cfg.LockThread),
		gnet.WithLoadBalancing(gnet.LeastConnections),
	)
}

type SockServer struct {
	gnet.BuiltinEventEngine

	ctx  context.Context
	cfg  Config
	pool *ants.Pool
	eng  gnet.Engine
}

func (mgr *SockServer) OnTraffic(conn gnet.Conn) gnet.Action {
	d, err := conn.Next(-1)
	if err != nil {
		log.Printf("[ERR] failed to read data from connection: %v\n", err)
		return gnet.Close
	}

	if mgr.cfg.AsyncWrite {
		if err := conn.AsyncWrite(d, nil); err != nil {
			log.Printf("[ERR] failed to write data to connection: %v\n", err)
			return gnet.Close
		}
	} else {
		if _, err := conn.Write(d); err != nil {
			log.Printf("[ERR] failed to write data to connection: %v\n", err)
			return gnet.Close
		}
	}

	return gnet.None
}

// OnBoot is called when the gnet engine is booted and ready to
// accept incoming connections.
func (mgr *SockServer) OnBoot(eng gnet.Engine) gnet.Action {
	mgr.eng = eng
	return gnet.None
}

// OnOpen is called when a new connection is opened.
func (mgr *SockServer) OnOpen(conn gnet.Conn) (out []byte, action gnet.Action) {
	return nil, gnet.None
}

func (mgr *SockServer) OnClose(conn gnet.Conn, err error) (action gnet.Action) {
	return gnet.None
}

func (mgr *SockServer) OnTick() (delay time.Duration, action gnet.Action) {
	select {
	case <-mgr.ctx.Done():
		log.Printf("server is exiting\n")
		return -1, gnet.Shutdown

	default:
		log.Printf("connections: %d\n", mgr.eng.CountConnections())
		return 1 * time.Second, gnet.None
	}
}