From 05e1758c00499f27c17eec50911845329bdd3994 Mon Sep 17 00:00:00 2001 From: youfak Date: Fri, 31 Jul 2026 11:34:01 +0800 Subject: [PATCH] feat: handle worker sessions and runtime reports --- internal/controller/worker/service.go | 192 +++++++++++++++++++++ internal/controller/worker/service_test.go | 101 +++++++++++ 2 files changed, 293 insertions(+) create mode 100644 internal/controller/worker/service.go create mode 100644 internal/controller/worker/service_test.go diff --git a/internal/controller/worker/service.go b/internal/controller/worker/service.go new file mode 100644 index 0000000..b51a396 --- /dev/null +++ b/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) +} diff --git a/internal/controller/worker/service_test.go b/internal/controller/worker/service_test.go new file mode 100644 index 0000000..f150d47 --- /dev/null +++ b/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") +}