proto-socket/go/inbound_gateway.go

package proto_socket

import (
	"runtime"
	"sync"

	"google.golang.org/protobuf/proto"

	"git.toki-labs.com/toki/proto-socket/go/packets"
)

// InboundFrame is a raw inbound frame queued for gateway preprocessing.
//
// Seq is an internal monotonically increasing sequence assigned by the receive
// side before submitting to the gateway. It is used only to restore input order
// after parallel decoding and never touches the wire format or any public
// protocol field.
type InboundFrame struct {
	Seq   int64
	Bytes []byte
}

// DecodedFrame is the result of decoding a PacketBase envelope inside a gateway.
//
// It carries the same envelope fields that Communicator.OnReceivedData expects,
// plus the originating Seq so results can be reordered before dispatch.
type DecodedFrame struct {
	Seq           int64
	TypeName      string
	Data          []byte
	IncomingNonce int32
	ResponseNonce int32
}

// InboundGateway performs only pure raw-bytes work (PacketBase envelope decode /
// lightweight preprocessing) off the receive coordinator. It must NOT own
// stateful dispatch, the pending-request map, listeners, or the outbound write
// queue — those stay on the receive coordinator. Decoded results are delivered
// to a sink in input Seq order so the coordinator can dispatch deterministically.
type InboundGateway interface {
	// Submit hands a raw frame to the gateway. Fire-and-forget; the decoded
	// result is later delivered to the sink in Seq order.
	Submit(frame InboundFrame)
	// Close releases gateway resources and stops delivering results.
	Close()
}

// decodeFunc decodes a raw frame into a PacketBase envelope. Pure: no shared
// state, safe to run concurrently across workers.
type decodeFunc func([]byte) (*packets.PacketBase, error)

func decodePacketBase(raw []byte) (*packets.PacketBase, error) {
	base := &packets.PacketBase{}
	if err := proto.Unmarshal(raw, base); err != nil {
		return nil, err
	}
	return base, nil
}

// workerResult is a per-frame decode outcome carried through the reorder buffer.
// err frames still advance the sequence so a single bad frame cannot stall the
// reorder window.
type workerResult struct {
	seq   int64
	frame DecodedFrame
	err   error
}

// frameReorderBuffer restores input order from results that may arrive out of
// order. Sequences are expected to be contiguous starting at 1. A result is
// buffered until every earlier sequence has been released, then it and any
// now-contiguous successors are returned in order. It is not safe for concurrent
// use; the WorkerGateway confines it to the single collector goroutine.
type frameReorderBuffer struct {
	nextSeq int64
	pending map[int64]workerResult
}

func newFrameReorderBuffer() *frameReorderBuffer {
	return &frameReorderBuffer{nextSeq: 1, pending: make(map[int64]workerResult)}
}

// release buffers res and returns the results now releasable in seq order.
func (b *frameReorderBuffer) release(res workerResult) []workerResult {
	if res.seq < b.nextSeq {
		// Already released (or duplicate); ignore to avoid emitting twice.
		return nil
	}
	b.pending[res.seq] = res
	var ready []workerResult
	for {
		r, ok := b.pending[b.nextSeq]
		if !ok {
			break
		}
		ready = append(ready, r)
		delete(b.pending, b.nextSeq)
		b.nextSeq++
	}
	return ready
}

// WorkerGateway is a goroutine worker-pool gateway. A bounded jobs channel feeds
// N workers that decode PacketBase envelopes in parallel; a single collector
// goroutine reorders the results by Seq and delivers them to the sink in input
// order. Because the collector is the only goroutine that calls the sink, sink
// invocations are serialized and ordered, keeping stateful dispatch on the
// receive coordinator.
type WorkerGateway struct {
	jobs          chan InboundFrame
	results       chan workerResult
	done          chan struct{}
	collectorDone chan struct{}
	wg            sync.WaitGroup
	closeOnce     sync.Once
	parse         decodeFunc
	sink          func(DecodedFrame)
	onError       func(error)
}

// NewWorkerGateway creates a worker-pool gateway that decodes PacketBase frames
// and delivers them to sink in Seq order. A frame that fails to decode is not
// sent to sink; instead onError is invoked with the decode error in Seq order so
// the transport can apply its existing parse-error disconnect semantics. The
// gateway itself never tears down the transport — onError keeps that decision on
// the caller. onError may be nil to ignore decode errors.
//
// workers <= 0 defaults to NumCPU and queueSize <= 0 defaults to 64. The bounded
// jobs channel provides backpressure: Submit blocks once it is full until a
// worker frees a slot.
func NewWorkerGateway(workers, queueSize int, sink func(DecodedFrame), onError func(error)) *WorkerGateway {
	return newWorkerGateway(workers, queueSize, decodePacketBase, sink, onError)
}

func newWorkerGateway(workers, queueSize int, parse decodeFunc, sink func(DecodedFrame), onError func(error)) *WorkerGateway {
	if workers <= 0 {
		workers = runtime.NumCPU()
	}
	if queueSize <= 0 {
		queueSize = 64
	}
	g := &WorkerGateway{
		jobs:          make(chan InboundFrame, queueSize),
		results:       make(chan workerResult, queueSize),
		done:          make(chan struct{}),
		collectorDone: make(chan struct{}),
		parse:         parse,
		sink:          sink,
		onError:       onError,
	}
	for i := 0; i < workers; i++ {
		g.wg.Add(1)
		go g.worker()
	}
	go g.collector()
	return g
}

func (g *WorkerGateway) worker() {
	defer g.wg.Done()
	for {
		select {
		case frame := <-g.jobs:
			res := workerResult{seq: frame.Seq}
			base, err := g.parse(frame.Bytes)
			if err != nil {
				res.err = err
			} else {
				res.frame = DecodedFrame{
					Seq:           frame.Seq,
					TypeName:      base.GetTypeName(),
					Data:          base.GetData(),
					IncomingNonce: base.GetNonce(),
					ResponseNonce: base.GetResponseNonce(),
				}
			}
			select {
			case g.results <- res:
			case <-g.done:
				return
			}
		case <-g.done:
			return
		}
	}
}

func (g *WorkerGateway) collector() {
	defer close(g.collectorDone)
	reorder := newFrameReorderBuffer()
	for {
		select {
		case res := <-g.results:
			for _, ready := range reorder.release(res) {
				if ready.err != nil {
					// Report the decode error in Seq order and keep going so a
					// single bad frame does not stall later frames. The caller's
					// onError decides whether to tear down the transport.
					if g.onError != nil {
						g.onError(ready.err)
					}
					continue
				}
				g.sink(ready.frame)
			}
		case <-g.done:
			return
		}
	}
}

// Submit hands a raw frame to the worker pool. It blocks while the bounded jobs
// channel is full (backpressure) and returns immediately once the gateway is
// closed.
func (g *WorkerGateway) Submit(frame InboundFrame) {
	select {
	case g.jobs <- frame:
	case <-g.done:
	}
}

// Close signals all workers and the collector to stop. It is idempotent and
// does not block on in-flight sink calls.
func (g *WorkerGateway) Close() {
	g.closeOnce.Do(func() {
		close(g.done)
	})
}