feat: add bounded session scheduling

2026-09-02 06:02:44 +00:00
parent 2764024d36
commit cfb2693930
3 changed files with 401 additions and 0 deletions
@@ -0,0 +1,277 @@
+// Package session owns transport-neutral server session scheduling and fencing.
+package session
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"time"
+)
+
+var (
+	ErrControlLaneFull = errors.New("session control lane is full")
+	ErrSessionClosed   = errors.New("session is closed")
+)
+
+type FrameKind uint8
+
+const (
+	FrameControl FrameKind = iota + 1
+	FrameData
+)
+
+type Frame struct {
+	Kind    FrameKind
+	Payload []byte
+}
+
+// WriterQueue is owned by one socket writer. Data saturation leaves the work
+// durable; control saturation is a session-fatal invariant.
+type WriterQueue struct {
+	control      chan Frame
+	data         chan Frame
+	controlBurst uint32
+	controlRun   uint32
+	closed       chan struct{}
+	closeOnce    sync.Once
+}
+
+func NewWriterQueue(controlCapacity, dataCapacity, controlBurst uint32) *WriterQueue {
+	if controlCapacity == 0 || dataCapacity == 0 || controlBurst == 0 {
+		panic("session queue capacities must be positive")
+	}
+	return &WriterQueue{
+		control: make(chan Frame, controlCapacity), data: make(chan Frame, dataCapacity),
+		controlBurst: controlBurst, closed: make(chan struct{}),
+	}
+}
+
+func (queue *WriterQueue) EnqueueControl(frame Frame) error {
+	if frame.Kind != FrameControl {
+		return ErrControlLaneFull
+	}
+	select {
+	case <-queue.closed:
+		return ErrSessionClosed
+	default:
+	}
+	select {
+	case queue.control <- frame:
+		return nil
+	default:
+		return ErrControlLaneFull
+	}
+}
+
+// EnqueueData reports false when the bounded data lane is full. It never
+// creates a fallback goroutine or buffer.
+func (queue *WriterQueue) EnqueueData(frame Frame) bool {
+	if frame.Kind != FrameData {
+		return false
+	}
+	select {
+	case <-queue.closed:
+		return false
+	default:
+	}
+	select {
+	case queue.data <- frame:
+		return true
+	default:
+		return false
+	}
+}
+
+func (queue *WriterQueue) Next(ctx context.Context) (Frame, error) {
+	for {
+		if queue.controlRun < queue.controlBurst {
+			select {
+			case frame := <-queue.control:
+				queue.controlRun++
+				return frame, nil
+			default:
+			}
+		}
+		select {
+		case frame := <-queue.data:
+			queue.controlRun = 0
+			return frame, nil
+		default:
+		}
+		select {
+		case frame := <-queue.control:
+			queue.controlRun++
+			return frame, nil
+		default:
+		}
+		select {
+		case <-ctx.Done():
+			return Frame{}, ctx.Err()
+		case <-queue.closed:
+			return Frame{}, ErrSessionClosed
+		case frame := <-queue.control:
+			queue.controlRun++
+			return frame, nil
+		case frame := <-queue.data:
+			queue.controlRun = 0
+			return frame, nil
+		}
+	}
+}
+
+func (queue *WriterQueue) Close() { queue.closeOnce.Do(func() { close(queue.closed) }) }
+
+type HeartbeatAction uint8
+
+const (
+	HeartbeatNone HeartbeatAction = iota
+	HeartbeatPing
+	HeartbeatClose
+)
+
+type Heartbeat struct {
+	idle        time.Duration
+	timeout     time.Duration
+	lastInbound time.Duration
+	pingSent    bool
+}
+
+func NewHeartbeat(idle, timeout, now time.Duration) *Heartbeat {
+	if idle <= 0 || timeout <= idle {
+		panic("invalid heartbeat timings")
+	}
+	return &Heartbeat{idle: idle, timeout: timeout, lastInbound: now}
+}
+
+func (heartbeat *Heartbeat) ObserveInbound(now time.Duration) {
+	if now >= heartbeat.lastInbound {
+		heartbeat.lastInbound = now
+		heartbeat.pingSent = false
+	}
+}
+
+func (heartbeat *Heartbeat) Check(now time.Duration) HeartbeatAction {
+	if now < heartbeat.lastInbound {
+		return HeartbeatNone
+	}
+	elapsed := now - heartbeat.lastInbound
+	if elapsed >= heartbeat.timeout {
+		return HeartbeatClose
+	}
+	if elapsed >= heartbeat.idle && !heartbeat.pingSent {
+		heartbeat.pingSent = true
+		return HeartbeatPing
+	}
+	return HeartbeatNone
+}
+
+type DispatchLane uint8
+
+const (
+	DispatchNone DispatchLane = iota
+	DispatchRunning
+	DispatchQueued
+)
+
+type CapacityShadow struct {
+	running, queued       uint32
+	maxRunning, maxQueued uint32
+	shadowRunning         uint32
+	shadowQueued          uint32
+}
+
+func (shadow *CapacityShadow) UpdateAdvertised(running, queued, maxRunning, maxQueued uint32) bool {
+	if maxRunning == 0 || maxQueued == 0 || running > maxRunning || queued > maxQueued {
+		return false
+	}
+	shadow.running, shadow.queued = running, queued
+	shadow.maxRunning, shadow.maxQueued = maxRunning, maxQueued
+	return true
+}
+
+func (shadow *CapacityShadow) Reserve() DispatchLane {
+	if shadow.running+shadow.shadowRunning < shadow.maxRunning {
+		shadow.shadowRunning++
+		return DispatchRunning
+	}
+	if shadow.queued+shadow.shadowQueued < shadow.maxQueued {
+		shadow.shadowQueued++
+		return DispatchQueued
+	}
+	return DispatchNone
+}
+
+func (shadow *CapacityShadow) Release(lane DispatchLane) bool {
+	switch lane {
+	case DispatchRunning:
+		if shadow.shadowRunning == 0 {
+			return false
+		}
+		shadow.shadowRunning--
+	case DispatchQueued:
+		if shadow.shadowQueued == 0 {
+			return false
+		}
+		shadow.shadowQueued--
+	default:
+		return false
+	}
+	return true
+}
+
+func (shadow CapacityShadow) Pending() (running, queued uint32) {
+	return shadow.shadowRunning, shadow.shadowQueued
+}
+
+type Handle struct {
+	ClientID   string
+	SessionID  [16]byte
+	Generation uint64
+	Context    context.Context
+	cancel     context.CancelFunc
+}
+
+type Registry struct {
+	mu      sync.Mutex
+	clients map[string]*Handle
+}
+
+func NewRegistry() *Registry { return &Registry{clients: make(map[string]*Handle)} }
+
+// Install replaces only an older in-memory handle for the same client. The
+// caller must have already durably fenced it in the store.
+func (registry *Registry) Install(clientID string, sessionID [16]byte, generation uint64) (*Handle, error) {
+	if clientID == "" || sessionID == [16]byte{} || generation == 0 {
+		return nil, ErrSessionClosed
+	}
+	registry.mu.Lock()
+	defer registry.mu.Unlock()
+	if old := registry.clients[clientID]; old != nil {
+		old.cancel()
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	handle := &Handle{ClientID: clientID, SessionID: sessionID, Generation: generation, Context: ctx, cancel: cancel}
+	registry.clients[clientID] = handle
+	return handle, nil
+}
+
+func (registry *Registry) Remove(handle *Handle) bool {
+	if handle == nil {
+		return false
+	}
+	registry.mu.Lock()
+	defer registry.mu.Unlock()
+	current := registry.clients[handle.ClientID]
+	if current != handle || current.Generation != handle.Generation || current.SessionID != handle.SessionID {
+		return false
+	}
+	delete(registry.clients, handle.ClientID)
+	handle.cancel()
+	return true
+}
+
+func (registry *Registry) Get(clientID string) *Handle {
+	registry.mu.Lock()
+	defer registry.mu.Unlock()
+	return registry.clients[clientID]
+}
@@ -0,0 +1,106 @@
+package session
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+)
+
+func TestWriterQueuePriorityAndFairness_HP_SES_03(t *testing.T) {
+	t.Parallel()
+	queue := NewWriterQueue(4, 1, 2)
+	if !queue.EnqueueData(Frame{Kind: FrameData, Payload: []byte("data")}) {
+		t.Fatal("initial data enqueue failed")
+	}
+	if queue.EnqueueData(Frame{Kind: FrameData}) {
+		t.Fatal("full data lane accepted another frame")
+	}
+	for _, value := range []string{"control-1", "control-2", "control-3"} {
+		if err := queue.EnqueueControl(Frame{Kind: FrameControl, Payload: []byte(value)}); err != nil {
+			t.Fatal(err)
+		}
+	}
+	ctx := context.Background()
+	for index, want := range []string{"control-1", "control-2", "data", "control-3"} {
+		frame, err := queue.Next(ctx)
+		if err != nil || string(frame.Payload) != want {
+			t.Fatalf("frame %d = (%q, %v), want %q", index, frame.Payload, err, want)
+		}
+	}
+	if err := queue.EnqueueControl(Frame{Kind: FrameData}); !errors.Is(err, ErrControlLaneFull) {
+		t.Fatalf("wrong control frame error = %v", err)
+	}
+	queue.Close()
+	if _, err := queue.Next(context.Background()); !errors.Is(err, ErrSessionClosed) {
+		t.Fatalf("closed queue Next error = %v", err)
+	}
+}
+
+func TestHeartbeatBoundaries_HP_SES_04(t *testing.T) {
+	t.Parallel()
+	heartbeat := NewHeartbeat(10*time.Second, 30*time.Second, 0)
+	if action := heartbeat.Check(9 * time.Second); action != HeartbeatNone {
+		t.Fatalf("9s action = %d", action)
+	}
+	if action := heartbeat.Check(10 * time.Second); action != HeartbeatPing {
+		t.Fatalf("10s action = %d", action)
+	}
+	if action := heartbeat.Check(29 * time.Second); action != HeartbeatNone {
+		t.Fatalf("29s action = %d", action)
+	}
+	if action := heartbeat.Check(30 * time.Second); action != HeartbeatClose {
+		t.Fatalf("30s action = %d", action)
+	}
+	heartbeat.ObserveInbound(31 * time.Second)
+	if action := heartbeat.Check(41 * time.Second); action != HeartbeatPing {
+		t.Fatalf("inbound reset action = %d", action)
+	}
+	if action := heartbeat.Check(1 * time.Second); action != HeartbeatNone {
+		t.Fatalf("clock rollback action = %d", action)
+	}
+}
+
+func TestCapacityShadowAndRegistry_BH_SES_03(t *testing.T) {
+	t.Parallel()
+	var shadow CapacityShadow
+	if shadow.UpdateAdvertised(0, 0, 0, 1) {
+		t.Fatal("invalid capacity accepted")
+	}
+	if !shadow.UpdateAdvertised(0, 0, 1, 1) {
+		t.Fatal("valid capacity rejected")
+	}
+	if got := shadow.Reserve(); got != DispatchRunning {
+		t.Fatalf("first lane = %d", got)
+	}
+	if got := shadow.Reserve(); got != DispatchQueued {
+		t.Fatalf("second lane = %d", got)
+	}
+	if got := shadow.Reserve(); got != DispatchNone {
+		t.Fatalf("oversubscription lane = %d", got)
+	}
+	if !shadow.Release(DispatchRunning) || shadow.Release(DispatchRunning) {
+		t.Fatal("shadow release accounting is wrong")
+	}
+
+	registry := NewRegistry()
+	first, err := registry.Install("client-a", [16]byte{1}, 1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	second, err := registry.Install("client-a", [16]byte{2}, 2)
+	if err != nil {
+		t.Fatal(err)
+	}
+	select {
+	case <-first.Context.Done():
+	default:
+		t.Fatal("replacement did not fence old handle")
+	}
+	if registry.Remove(first) || registry.Get("client-a") != second {
+		t.Fatal("stale handle removed current session")
+	}
+	if !registry.Remove(second) {
+		t.Fatal("current handle removal failed")
+	}
+}