feat: handle worker sessions and runtime reports
This commit is contained in:
youfak
parent
a463a8cbd2
commit
05e1758c00
192
internal/controller/worker/service.go
Normal file
192
internal/controller/worker/service.go
Normal file
@ -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)
|
||||
}
|
||||
101
internal/controller/worker/service_test.go
Normal file
101
internal/controller/worker/service_test.go
Normal file
@ -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")
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user