feat: handle worker sessions and runtime reports

internal/controller/worker/service.go

@@ -0,0 +1,192 @@
+package worker
+
+import (
+	"context"
+	"crypto/rand"
+	"crypto/sha256"
+	"encoding/hex"
+	"errors"
+	"time"
+
+	"proxy-pool/internal/domain/workerruntime"
+)
+
+var (
+	ErrInvalidCommand  = errors.New("invalid worker control command")
+	ErrProtocolVersion = errors.New("unsupported worker protocol version")
+	ErrUnavailable     = errors.New("worker control service unavailable")
+)
+
+type RegisterCommand struct {
+	WorkerID        string
+	InstanceID      string
+	Zone            string
+	ProtocolVersion uint32
+	Labels          map[string]string
+}
+
+type SnapshotAcknowledgement struct {
+	WorkerID       string
+	SessionID      string
+	Version        uint64
+	OwnershipEpoch uint64
+	Checksum       []byte
+	Applied        bool
+	ErrorCode      string
+	ErrorMessage   string
+}
+
+type Registration struct {
+	WorkerID          string
+	SessionID         string
+	OwnershipEpoch    uint64
+	HeartbeatInterval time.Duration
+	MaxStaleAge       time.Duration
+}
+
+type RuntimeDecision struct {
+	AcceptedOwnershipEpoch uint64
+	RequireFullSnapshot    bool
+}
+
+type Options struct {
+	ProtocolVersion    uint32
+	HeartbeatInterval  time.Duration
+	SessionTTL         time.Duration
+	MaxStaleAge        time.Duration
+	MaxRuntimeCounters int
+	SessionID          func() (string, error)
+}
+
+type Service interface {
+	Register(context.Context, RegisterCommand) (Registration, error)
+	Acknowledge(context.Context, SnapshotAcknowledgement) error
+	ReportRuntime(context.Context, workerruntime.Report) (RuntimeDecision, error)
+}
+
+type service struct {
+	store   workerruntime.ControlStore
+	options Options
+}
+
+func NewService(store workerruntime.ControlStore, options Options) (Service, error) {
+	if store == nil || options.ProtocolVersion == 0 || options.HeartbeatInterval <= 0 ||
+		options.SessionTTL < 3*options.HeartbeatInterval || options.MaxStaleAge < options.HeartbeatInterval ||
+		options.MaxRuntimeCounters <= 0 {
+		return nil, ErrInvalidCommand
+	}
+	if options.SessionID == nil {
+		options.SessionID = randomSessionID
+	}
+	return &service{store: store, options: options}, nil
+}
+
+func (service *service) Register(ctx context.Context, command RegisterCommand) (Registration, error) {
+	if ctx == nil {
+		return Registration{}, ErrInvalidCommand
+	}
+	if err := ctx.Err(); err != nil {
+		return Registration{}, err
+	}
+	if command.ProtocolVersion != service.options.ProtocolVersion {
+		return Registration{}, ErrProtocolVersion
+	}
+	epoch, err := service.store.CurrentOwnershipEpoch(ctx)
+	if err != nil {
+		return Registration{}, classifyStoreError(err)
+	}
+	sessionID, err := service.options.SessionID()
+	if err != nil || !workerruntime.ValidIdentifier(sessionID) {
+		return Registration{}, errors.Join(ErrUnavailable, err)
+	}
+	session := workerruntime.Session{
+		WorkerID: command.WorkerID, InstanceID: command.InstanceID, SessionID: sessionID,
+		Zone: command.Zone, ProtocolVersion: command.ProtocolVersion, Labels: command.Labels,
+	}
+	if _, err := workerruntime.NormalizeSession(session); err != nil {
+		return Registration{}, errors.Join(ErrInvalidCommand, err)
+	}
+	if err := service.store.OpenSession(ctx, session, service.options.SessionTTL); err != nil {
+		return Registration{}, classifyStoreError(err)
+	}
+	return Registration{
+		WorkerID: command.WorkerID, SessionID: sessionID, OwnershipEpoch: epoch,
+		HeartbeatInterval: service.options.HeartbeatInterval, MaxStaleAge: service.options.MaxStaleAge,
+	}, nil
+}
+
+func (service *service) Acknowledge(ctx context.Context, acknowledgement SnapshotAcknowledgement) error {
+	if ctx == nil || len(acknowledgement.Checksum) != sha256.Size || len(acknowledgement.ErrorMessage) > 512 {
+		return ErrInvalidCommand
+	}
+	if err := ctx.Err(); err != nil {
+		return err
+	}
+	var checksum [sha256.Size]byte
+	copy(checksum[:], acknowledgement.Checksum)
+	domainAcknowledgement := workerruntime.SnapshotAcknowledgement{
+		WorkerID: acknowledgement.WorkerID, SessionID: acknowledgement.SessionID,
+		Reference: workerruntime.SnapshotReference{
+			WorkerID: acknowledgement.WorkerID, Version: acknowledgement.Version,
+			OwnershipEpoch: acknowledgement.OwnershipEpoch, Checksum: checksum,
+		},
+		Applied: acknowledgement.Applied, ErrorCode: acknowledgement.ErrorCode,
+	}
+	if _, err := workerruntime.NormalizeAcknowledgement(domainAcknowledgement); err != nil {
+		return errors.Join(ErrInvalidCommand, err)
+	}
+	if err := service.store.AcknowledgeSnapshot(ctx, domainAcknowledgement, service.options.SessionTTL); err != nil {
+		return classifyStoreError(err)
+	}
+	return nil
+}
+
+func (service *service) ReportRuntime(ctx context.Context, report workerruntime.Report) (RuntimeDecision, error) {
+	if ctx == nil {
+		return RuntimeDecision{}, ErrInvalidCommand
+	}
+	if err := ctx.Err(); err != nil {
+		return RuntimeDecision{}, err
+	}
+	normalized, _, err := workerruntime.NormalizeReport(report)
+	if err != nil || len(normalized.Counters) > service.options.MaxRuntimeCounters {
+		return RuntimeDecision{}, errors.Join(ErrInvalidCommand, err)
+	}
+	if err := service.store.ReplaceRuntime(ctx, normalized, service.options.SessionTTL); err != nil {
+		if errors.Is(err, workerruntime.ErrSnapshotMismatch) {
+			epoch, epochErr := service.store.CurrentOwnershipEpoch(ctx)
+			if epochErr != nil {
+				return RuntimeDecision{}, classifyStoreError(epochErr)
+			}
+			return RuntimeDecision{AcceptedOwnershipEpoch: epoch, RequireFullSnapshot: true}, nil
+		}
+		return RuntimeDecision{}, classifyStoreError(err)
+	}
+	epoch, err := service.store.CurrentOwnershipEpoch(ctx)
+	if err != nil {
+		return RuntimeDecision{}, classifyStoreError(err)
+	}
+	return RuntimeDecision{AcceptedOwnershipEpoch: epoch}, nil
+}
+
+func randomSessionID() (string, error) {
+	var value [16]byte
+	if _, err := rand.Read(value[:]); err != nil {
+		return "", err
+	}
+	return hex.EncodeToString(value[:]), nil
+}
+
+func classifyStoreError(err error) error {
+	if err == nil {
+		return nil
+	}
+	if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) ||
+		errors.Is(err, workerruntime.ErrInvalidSession) || errors.Is(err, workerruntime.ErrInvalidReport) ||
+		errors.Is(err, workerruntime.ErrInvalidAcknowledgement) || errors.Is(err, workerruntime.ErrStaleSession) ||
+		errors.Is(err, workerruntime.ErrStaleReport) || errors.Is(err, workerruntime.ErrConflictingReport) ||
+		errors.Is(err, workerruntime.ErrSnapshotMismatch) || errors.Is(err, workerruntime.ErrStaleAcknowledgement) {
+		return err
+	}
+	return errors.Join(ErrUnavailable, err)
+}

internal/controller/worker/service_test.go

@@ -0,0 +1,101 @@
+package worker
+
+import (
+	"context"
+	"crypto/sha256"
+	"errors"
+	"testing"
+	"time"
+
+	"proxy-pool/internal/domain/workerruntime"
+)
+
+func TestServiceRegistersAcknowledgesAndReportsRuntime(t *testing.T) {
+	now := time.Date(2026, 7, 31, 10, 0, 0, 0, time.UTC)
+	store, err := workerruntime.NewMemoryStore(func() time.Time { return now })
+	if err != nil {
+		t.Fatalf("NewMemoryStore(): %v", err)
+	}
+	service, err := NewService(store, Options{
+		ProtocolVersion: 1, HeartbeatInterval: 10 * time.Second, SessionTTL: 30 * time.Second,
+		MaxStaleAge: 10 * time.Second, MaxRuntimeCounters: 100,
+		SessionID: func() (string, error) { return "0123456789abcdef0123456789abcdef", nil },
+	})
+	if err != nil {
+		t.Fatalf("NewService(): %v", err)
+	}
+	registered, err := service.Register(context.Background(), RegisterCommand{
+		WorkerID: "worker-a", InstanceID: "instance-a", Zone: "zone-a", ProtocolVersion: 1,
+		Labels: map[string]string{"region": "test"},
+	})
+	if err != nil || registered.SessionID == "" || registered.OwnershipEpoch == 0 {
+		t.Fatalf("Register() = %+v, %v", registered, err)
+	}
+	reference := workerruntime.SnapshotReference{
+		WorkerID: "worker-a", Version: 7, OwnershipEpoch: registered.OwnershipEpoch,
+		Checksum: sha256.Sum256([]byte("snapshot-7")),
+	}
+	if err := store.RecordIssuedSnapshot(context.Background(), reference, time.Minute); err != nil {
+		t.Fatalf("RecordIssuedSnapshot(): %v", err)
+	}
+	if err := service.Acknowledge(context.Background(), SnapshotAcknowledgement{
+		WorkerID: "worker-a", SessionID: registered.SessionID, Version: 7,
+		OwnershipEpoch: registered.OwnershipEpoch, Checksum: reference.Checksum[:], Applied: true,
+	}); err != nil {
+		t.Fatalf("Acknowledge(): %v", err)
+	}
+	decision, err := service.ReportRuntime(context.Background(), workerruntime.Report{
+		WorkerID: "worker-a", SessionID: registered.SessionID, Sequence: 1,
+		SnapshotVersion: 7, OwnershipEpoch: registered.OwnershipEpoch, ObservedAt: now,
+	})
+	if err != nil || decision.RequireFullSnapshot || decision.AcceptedOwnershipEpoch != registered.OwnershipEpoch {
+		t.Fatalf("ReportRuntime() = %+v, %v", decision, err)
+	}
+	decision, err = service.ReportRuntime(context.Background(), workerruntime.Report{
+		WorkerID: "worker-a", SessionID: registered.SessionID, Sequence: 2,
+		SnapshotVersion: 8, OwnershipEpoch: registered.OwnershipEpoch, ObservedAt: now,
+	})
+	if err != nil || !decision.RequireFullSnapshot || decision.AcceptedOwnershipEpoch != registered.OwnershipEpoch {
+		t.Fatalf("ReportRuntime(mismatch) = %+v, %v", decision, err)
+	}
+}
+
+func TestServiceRejectsInvalidCommandsAndUnavailableStore(t *testing.T) {
+	store := unavailableStore{}
+	service, err := NewService(store, Options{
+		ProtocolVersion: 1, HeartbeatInterval: time.Second, SessionTTL: 3 * time.Second,
+		MaxStaleAge: time.Second, MaxRuntimeCounters: 1,
+	})
+	if err != nil {
+		t.Fatalf("NewService(): %v", err)
+	}
+	if _, err := service.Register(context.Background(), RegisterCommand{ProtocolVersion: 2}); !errors.Is(err, ErrProtocolVersion) {
+		t.Fatalf("Register(protocol) error = %v", err)
+	}
+	if _, err := service.Register(context.Background(), RegisterCommand{
+		WorkerID: "worker-a", InstanceID: "instance-a", Zone: "zone-a", ProtocolVersion: 1,
+	}); !errors.Is(err, ErrUnavailable) {
+		t.Fatalf("Register(unavailable) error = %v", err)
+	}
+	if err := service.Acknowledge(context.Background(), SnapshotAcknowledgement{Checksum: make([]byte, 31)}); !errors.Is(err, ErrInvalidCommand) {
+		t.Fatalf("Acknowledge(invalid) error = %v", err)
+	}
+}
+
+type unavailableStore struct{}
+
+func (unavailableStore) CurrentOwnershipEpoch(context.Context) (uint64, error) {
+	return 0, errors.New("redis unavailable")
+}
+func (unavailableStore) OpenSession(context.Context, workerruntime.Session, time.Duration) error {
+	return errors.New("redis unavailable")
+}
+func (unavailableStore) RecordIssuedSnapshot(context.Context, workerruntime.SnapshotReference, time.Duration) error {
+	return errors.New("redis unavailable")
+}
+func (unavailableStore) AcknowledgeSnapshot(context.Context, workerruntime.SnapshotAcknowledgement, time.Duration) error {
+	return errors.New("redis unavailable")
+}
+func (unavailableStore) ReplaceRuntime(context.Context, workerruntime.Report, time.Duration) error {
+	return errors.New("redis unavailable")
+}