lingniu-vehicle-ingest/go/vehicle-gateway/internal/gateway/tcp_server.go

package gateway

import (
	"context"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net"
	"strings"
	"sync"
	"time"

	"lingniu-vehicle-ingest/go/vehicle-gateway/internal/envelope"
	"lingniu-vehicle-ingest/go/vehicle-gateway/internal/eventbus"
	"lingniu-vehicle-ingest/go/vehicle-gateway/internal/identity"
)

type FrameExtractor func([]byte) (frames [][]byte, remainder []byte, err error)

type FrameParser func(raw []byte, receivedAtMS int64, sourceEndpoint string) (envelope.FrameEnvelope, error)

type FrameResponder func(raw []byte, env envelope.FrameEnvelope) (response []byte, ok bool, err error)

type TCPProtocol struct {
	Protocol envelope.Protocol
	Addr     string
	Extract  FrameExtractor
	Parse    FrameParser
	Respond  FrameResponder
}

type TCPServer struct {
	protocol       TCPProtocol
	sink           eventbus.Sink
	resolver       identity.Resolver
	logger         *slog.Logger
	readBufferSize int
	idleTimeout    time.Duration
	maxConnections int
}

type TCPServerConfig struct {
	Protocol       TCPProtocol
	Sink           eventbus.Sink
	Resolver       identity.Resolver
	Logger         *slog.Logger
	ReadBufferSize int
	IdleTimeout    time.Duration
	MaxConnections int
}

func NewTCPServer(cfg TCPServerConfig) (*TCPServer, error) {
	if cfg.Protocol.Protocol == "" {
		return nil, errors.New("protocol is required")
	}
	if strings.TrimSpace(cfg.Protocol.Addr) == "" {
		return nil, errors.New("listen addr is required")
	}
	if cfg.Protocol.Extract == nil {
		return nil, errors.New("frame extractor is required")
	}
	if cfg.Protocol.Parse == nil {
		return nil, errors.New("frame parser is required")
	}
	if cfg.Sink == nil {
		return nil, errors.New("sink is required")
	}
	if cfg.Logger == nil {
		cfg.Logger = slog.Default()
	}
	if cfg.Resolver == nil {
		cfg.Resolver = identity.NoopResolver{}
	}
	if cfg.ReadBufferSize <= 0 {
		cfg.ReadBufferSize = 32 * 1024
	}
	if cfg.IdleTimeout <= 0 {
		cfg.IdleTimeout = 2 * time.Minute
	}
	if cfg.MaxConnections <= 0 {
		cfg.MaxConnections = 10_000
	}
	return &TCPServer{
		protocol:       cfg.Protocol,
		sink:           cfg.Sink,
		resolver:       cfg.Resolver,
		logger:         cfg.Logger,
		readBufferSize: cfg.ReadBufferSize,
		idleTimeout:    cfg.IdleTimeout,
		maxConnections: cfg.MaxConnections,
	}, nil
}

func (s *TCPServer) ListenAndServe(ctx context.Context) error {
	var lc net.ListenConfig
	listener, err := lc.Listen(ctx, "tcp", s.protocol.Addr)
	if err != nil {
		return err
	}
	defer listener.Close()

	go func() {
		<-ctx.Done()
		_ = listener.Close()
	}()

	s.logger.Info("tcp listener started", "protocol", s.protocol.Protocol, "addr", listener.Addr().String())
	sem := make(chan struct{}, s.maxConnections)
	var wg sync.WaitGroup
	defer wg.Wait()

	for {
		conn, err := listener.Accept()
		if err != nil {
			if ctx.Err() != nil {
				return nil
			}
			s.logger.Warn("tcp accept failed", "protocol", s.protocol.Protocol, "error", err)
			continue
		}
		select {
		case sem <- struct{}{}:
			wg.Add(1)
			go func() {
				defer wg.Done()
				defer func() { <-sem }()
				s.handleConnection(ctx, conn)
			}()
		default:
			s.logger.Warn("tcp connection rejected: max connections reached", "protocol", s.protocol.Protocol, "remote", conn.RemoteAddr().String())
			_ = conn.Close()
		}
	}
}

func (s *TCPServer) handleConnection(ctx context.Context, conn net.Conn) {
	defer conn.Close()
	source := conn.RemoteAddr().String()
	log := s.logger.With("protocol", s.protocol.Protocol, "remote", source)
	log.Info("tcp connection opened")
	defer log.Info("tcp connection closed")

	readBuffer := make([]byte, s.readBufferSize)
	var pending []byte
	for {
		_ = conn.SetReadDeadline(time.Now().Add(s.idleTimeout))
		n, err := conn.Read(readBuffer)
		if n > 0 {
			pending = append(pending, readBuffer[:n]...)
			frames, remainder, extractErr := s.protocol.Extract(pending)
			if extractErr != nil {
				log.Warn("frame extraction failed", "error", extractErr)
				return
			}
			pending = remainder
			for _, frame := range frames {
				s.handleFrame(ctx, conn, frame, source)
			}
		}
		if err != nil {
			if errors.Is(err, io.EOF) {
				return
			}
			var netErr net.Error
			if errors.As(err, &netErr) && netErr.Timeout() {
				log.Warn("tcp connection idle timeout")
				return
			}
			log.Warn("tcp read failed", "error", err)
			return
		}
		if ctx.Err() != nil {
			return
		}
	}
}

func (s *TCPServer) handleFrame(ctx context.Context, conn net.Conn, raw []byte, source string) {
	receivedAtMS := time.Now().UnixMilli()
	env, err := s.protocol.Parse(raw, receivedAtMS, source)
	if err != nil {
		env = envelope.FrameEnvelope{
			Protocol:       s.protocol.Protocol,
			SourceEndpoint: source,
			ReceivedAtMS:   receivedAtMS,
			EventTimeMS:    receivedAtMS,
			RawHex:         strings.ToUpper(hex.EncodeToString(raw)),
			ParseStatus:    envelope.ParseBadFrame,
			ParseError:     err.Error(),
		}
		env.EventID = env.StableEventID()
	} else {
		resolved, resolveErr := s.resolver.Resolve(ctx, env)
		if resolveErr != nil {
			s.logger.Warn("identity resolve failed", "protocol", s.protocol.Protocol, "event_id", env.StableEventID(), "error", resolveErr)
			if env.Parsed == nil {
				env.Parsed = map[string]any{}
			}
			env.Parsed["identity"] = map[string]any{"resolved": false, "error": resolveErr.Error()}
			env.ParseStatus = envelope.ParsePartial
		} else {
			env = resolved
		}
	}

	if err := s.sink.PublishRaw(ctx, env); err != nil {
		s.logger.Error("publish raw failed", "protocol", s.protocol.Protocol, "event_id", env.StableEventID(), "error", err)
		return
	}
	if env.ParseStatus == envelope.ParseBadFrame {
		return
	}
	if err := s.sink.PublishUnified(ctx, env); err != nil {
		s.logger.Error("publish unified failed", "protocol", s.protocol.Protocol, "event_id", env.StableEventID(), "error", err)
		return
	}
	if s.protocol.Respond == nil {
		return
	}
	response, ok, err := s.protocol.Respond(raw, env)
	if err != nil {
		s.logger.Warn("build protocol response failed", "protocol", s.protocol.Protocol, "event_id", env.StableEventID(), "error", err)
		return
	}
	if !ok || len(response) == 0 {
		return
	}
	_ = conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
	if _, err := conn.Write(response); err != nil {
		s.logger.Warn("write protocol response failed", "protocol", s.protocol.Protocol, "event_id", env.StableEventID(), "error", err)
	}
}

func (p TCPProtocol) String() string {
	return fmt.Sprintf("%s@%s", p.Protocol, p.Addr)
}