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

package session

import (
	"sync/atomic"
	"time"

	"spark-server/internal/pkg/errs"
	"spark-server/internal/pkg/reqctx"
	"spark-server/internal/transport/tcp/protocol"
)

const (
	StateConnecting uint32 = iota
	StateConnected
	StateAuthenticated
	StateActive
	StateClosing
	StateClosed
)

type Conn struct {
	id         string
	userID     atomic.Int64
	deviceID   atomic.Uint64
	realIP     atomic.Value // string
	language   atomic.Value // string
	remote     string
	state      atomic.Uint32
	lastActive atomic.Int64
	drops      atomic.Uint64
	proxyDone  atomic.Bool
	closed     atomic.Bool

	requestID string
	outbound chan *Packet
	write    func([]byte) error
	closeFn  func() error
	wake     func()
}

func NewConn(id, remote string, outboundSize int) *Conn {
	if outboundSize < 1 {
		outboundSize = 1024
	}
	c := &Conn{
		id:       id,
		remote:   remote,
		outbound: make(chan *Packet, outboundSize),
	}
	c.state.Store(StateConnected)
	c.lastActive.Store(time.Now().UnixNano())
	c.realIP.Store("")
	c.language.Store("")
	return c
}

func (c *Conn) SetIO(write func([]byte) error, closeFn func() error) {
	c.write = write
	c.closeFn = closeFn
}

// SetWake registers EventLoop wakeup after outbound enqueue.
func (c *Conn) SetWake(wake func()) {
	c.wake = wake
}

func (c *Conn) SetRequestID(id string) { c.requestID = id }
func (c *Conn) RequestID() string      { return c.requestID }

func (c *Conn) ID() string         { return c.id }
func (c *Conn) UserID() int64      { return c.userID.Load() }
func (c *Conn) DeviceID() uint64   { return c.deviceID.Load() }
func (c *Conn) RemoteAddr() string { return c.remote }
func (c *Conn) Drops() uint64      { return c.drops.Load() }

func (c *Conn) RealIP() string {
	v, _ := c.realIP.Load().(string)
	if v != "" {
		return v
	}
	return c.remote
}

func (c *Conn) SetRealIP(ip string) {
	c.realIP.Store(ip)
	c.proxyDone.Store(true)
}

func (c *Conn) Language() string {
	v, _ := c.language.Load().(string)
	return v
}

func (c *Conn) SetLanguage(lang string) {
	c.language.Store(lang)
}

func (c *Conn) MarkProxySkipped() { c.proxyDone.Store(true) }
func (c *Conn) ProxyDone() bool   { return c.proxyDone.Load() }

func (c *Conn) Touch() { c.lastActive.Store(time.Now().UnixNano()) }

func (c *Conn) Idle(timeout time.Duration) bool {
	if timeout <= 0 {
		return false
	}
	return time.Since(time.Unix(0, c.lastActive.Load())) > timeout
}

func (c *Conn) State() uint32 { return c.state.Load() }

func (c *Conn) Authenticated() bool {
	st := c.state.Load()
	return st == StateAuthenticated || st == StateActive
}

func (c *Conn) Authenticate(userID int64, deviceID uint64) {
	c.userID.Store(userID)
	c.deviceID.Store(deviceID)
	c.state.Store(StateAuthenticated)
}

func (c *Conn) Activate() {
	if c.state.Load() == StateAuthenticated {
		c.state.Store(StateActive)
	}
}

func (c *Conn) Send(packet *Packet) error {
	if c.closed.Load() {
		return errs.ErrClosed
	}
	select {
	case c.outbound <- packet:
		c.notifyWake()
		return nil
	default:
		c.drops.Add(1)
		return errs.ErrOutboundDropped
	}
}

func (c *Conn) notifyWake() {
	if c.wake != nil {
		c.wake()
	}
}

func (c *Conn) Close() error {
	if !c.closed.CompareAndSwap(false, true) {
		return nil
	}
	c.state.Store(StateClosing)
	c.flushDirect()
	c.state.Store(StateClosed)
	if c.closeFn != nil {
		return c.closeFn()
	}
	return nil
}

// Detach marks the session closed without touching the socket (OnClose).
func (c *Conn) Detach() {
	c.closed.Store(true)
	c.state.Store(StateClosed)
	c.write = nil
	c.closeFn = nil
	c.wake = nil
}

func (c *Conn) flushDirect() {
	for _, p := range c.DrainOutbound(64) {
		c.writePacket(p)
	}
}

func (c *Conn) writePacket(p *Packet) {
	if c.write == nil || p == nil {
		return
	}
	raw, err := protocol.Encode(protocol.Frame{
		Magic: protocol.Magic, Version: protocol.CurrentVersion,
		Command: p.Command, Flags: p.Flags, RequestID: p.RequestID, SequenceID: p.SequenceID, Body: p.Body,
	})
	if err != nil {
		return
	}
	_ = c.write(raw)
}

func (c *Conn) DrainOutbound(limit int) []*Packet {
	if limit < 1 {
		limit = 32
	}
	out := make([]*Packet, 0, limit)
	for len(out) < limit {
		select {
		case p := <-c.outbound:
			if p != nil {
				out = append(out, p)
			}
		default:
			return out
		}
	}
	return out
}

func (c *Conn) RequestContext() *reqctx.RequestContext {
	return &reqctx.RequestContext{
		RequestID:    c.requestID,
		UserID:       c.UserID(),
		DeviceID:     c.DeviceID(),
		ConnectionID: c.id,
		RemoteAddr:   c.remote,
		RealIP:       c.RealIP(),
		Language:     c.Language(),
	}
}