spark-server/internal/transport/tcp/server.go

package tcp

import (
	"context"
	"crypto/tls"
	"log/slog"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/panjf2000/gnet/v2"

	"spark-server/internal/module/auth"
	"spark-server/internal/transport/tcp/session"
)

const (
	defaultMaxFrame       = 65536
	defaultCompressMin    = 512
	defaultIdleTimeout    = 120 * time.Second
	defaultOutbound       = 1024
	defaultRatePerSec     = 20
	defaultHandlerTimeout = 5 * time.Second
	maxFramesPerEvent     = 32
)

// Metrics is a narrow counter hook. EventLoop must not block on exporters.
type Metrics interface {
	Inc(name string)
}

type nopMetrics struct{}

func (nopMetrics) Inc(string) {}

// MsgIDGenerator is satisfied by infrastructure/id.Generator.
type MsgIDGenerator interface {
	NextMsgID() (int64, error)
}

// Server is the gnet event engine. OnTraffic must not call MySQL, Redis, or deliver fanout.
type Server struct {
	gnet.BuiltinEventEngine
	log            *slog.Logger
	pool           TaskPool
	sessions       session.Manager
	eng            gnet.Engine
	addr           string
	maxFrame       int
	opts           []gnet.Option
	proxyProtocol  bool
	idleTimeout    time.Duration
	outboundSize   int
	ratePerSec     int
	compressMin    int
	handlerTimeout time.Duration
	auth           auth.Service
	ids            MsgIDGenerator
	router         *Router
	metrics        Metrics
	conns          sync.Map
	seq            atomic.Uint64
	ready          chan struct{}
	readyOnce      sync.Once
}

// Options configure the TCP listener. TLS is terminated at Nginx in production.
type Options struct {
	Addr           string
	Multicore      bool
	KeepAlive      time.Duration
	RecvBuffer     int
	SendBuffer     int
	MaxFrameSize   int
	TLS            *tls.Config
	ProxyProtocol  bool
	IdleTimeout    time.Duration
	OutboundSize   int
	RatePerSec     int
	CompressMin    int
	HandlerTimeout time.Duration
	Auth           auth.Service
	IDs            MsgIDGenerator
	Metrics        Metrics
}

// NewServer builds a gnet IM listener. Call Start to bind.
func NewServer(log *slog.Logger, pool TaskPool, sessions session.Manager, o *Options) *Server {
	if o == nil {
		o = &Options{}
	}
	applyTCPDefaults(o)
	s := &Server{
		log:            log,
		pool:           pool,
		sessions:       sessions,
		addr:           o.Addr,
		maxFrame:       o.MaxFrameSize,
		proxyProtocol:  o.ProxyProtocol,
		idleTimeout:    o.IdleTimeout,
		outboundSize:   o.OutboundSize,
		ratePerSec:     o.RatePerSec,
		compressMin:    o.CompressMin,
		handlerTimeout: o.HandlerTimeout,
		auth:           o.Auth,
		ids:            o.IDs,
		metrics:        o.Metrics,
		router:         NewRouter(),
		ready:          make(chan struct{}),
	}
	if s.metrics == nil {
		s.metrics = nopMetrics{}
	}
	s.registerDefaultHandlers()
	if o.TLS != nil {
		log.Info("tls cert loaded; terminate at nginx or wrap listener (gnet v2 has no native TLS)")
	}
	s.opts = gnetOptions(o)
	return s
}

func applyTCPDefaults(o *Options) {
	if o.MaxFrameSize <= 0 {
		o.MaxFrameSize = defaultMaxFrame
	}
	if o.CompressMin <= 0 {
		o.CompressMin = defaultCompressMin
	}
	if o.IdleTimeout <= 0 {
		o.IdleTimeout = defaultIdleTimeout
	}
	if o.OutboundSize <= 0 {
		o.OutboundSize = defaultOutbound
	}
	if o.RatePerSec <= 0 {
		o.RatePerSec = defaultRatePerSec
	}
	if o.HandlerTimeout <= 0 {
		o.HandlerTimeout = defaultHandlerTimeout
	}
}

func gnetOptions(o *Options) []gnet.Option {
	opts := []gnet.Option{
		gnet.WithMulticore(o.Multicore),
		gnet.WithReusePort(true),
		gnet.WithTicker(true),
	}
	if o.KeepAlive > 0 {
		opts = append(opts, gnet.WithTCPKeepAlive(o.KeepAlive))
	}
	if o.RecvBuffer > 0 {
		opts = append(opts, gnet.WithSocketRecvBuffer(o.RecvBuffer))
	}
	if o.SendBuffer > 0 {
		opts = append(opts, gnet.WithSocketSendBuffer(o.SendBuffer))
	}
	return opts
}

// Router returns the command table so modules can register SEND/SYNC later.
func (s *Server) Router() *Router { return s.router }

func (s *Server) OnBoot(eng gnet.Engine) gnet.Action {
	s.eng = eng
	s.log.Info("tcp listening", "addr", s.addr)
	s.readyOnce.Do(func() { close(s.ready) })
	return gnet.None
}

func (s *Server) OnOpen(c gnet.Conn) ([]byte, gnet.Action) {
	s.attach(c)
	return nil, gnet.None
}

func (s *Server) OnClose(c gnet.Conn, _ error) gnet.Action {
	st := connFrom(c)
	if st == nil {
		return gnet.None
	}
	s.conns.Delete(st.sess.ID())
	_ = s.sessions.Unbind(st.sess.ID())
	st.sess.Detach()
	return gnet.None
}

func (s *Server) OnTick() (time.Duration, gnet.Action) {
	s.conns.Range(func(_, v any) bool {
		st, ok := v.(*connState)
		if !ok {
			return true
		}
		if st.sess.Idle(s.idleTimeout) {
			_ = st.gc.Close()
			return true
		}
		s.flush(st)
		return true
	})
	return time.Second, gnet.None
}

func (s *Server) OnTraffic(c gnet.Conn) gnet.Action {
	st := connFrom(c)
	if st == nil {
		return gnet.Close
	}
	if action := s.consumeProxy(st); action != gnet.None {
		return action
	}
	if !st.sess.ProxyDone() {
		s.flush(st)
		return gnet.None
	}
	for i := 0; i < maxFramesPerEvent; i++ {
		action, more := s.consumeFrame(st)
		if action != gnet.None {
			return action
		}
		if !more {
			break
		}
	}
	s.flush(st)
	return gnet.None
}

func (s *Server) Start() error {
	return gnet.Run(s, "tcp://"+s.addr, s.opts...)
}

func (s *Server) Stop(ctx context.Context) error {
	return s.eng.Stop(ctx)
}

func (s *Server) attach(c gnet.Conn) {
	id := strconv.FormatUint(s.seq.Add(1), 10)
	remote := "unknown"
	if addr := c.RemoteAddr(); addr != nil {
		remote = addr.String()
	}
	sess := session.NewConn(id, remote, s.outboundSize)
	if !s.proxyProtocol {
		sess.MarkProxySkipped()
	}
	st := &connState{sess: sess, gc: c, limit: newLimiter(s.ratePerSec)}
	sess.SetIO(func(b []byte) error { return c.AsyncWrite(b, nil) }, func() error { return c.Close() })
	sess.SetWake(func() { _ = c.Wake(nil) })
	c.SetContext(st)
	s.conns.Store(id, st)
}