261 lines
8.3 KiB
Go
261 lines
8.3 KiB
Go
package worker
|
|
|
|
import (
|
|
"context"
|
|
"crypto/rand"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"errors"
|
|
"time"
|
|
|
|
ownershipDomain "proxy-pool/internal/domain/ownership"
|
|
"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
|
|
MaxSnapshotBytes int
|
|
SnapshotReader ownershipDomain.SnapshotReader
|
|
SessionID func() (string, error)
|
|
}
|
|
|
|
type Service interface {
|
|
Register(context.Context, RegisterCommand) (Registration, error)
|
|
CurrentOwnershipEpoch(context.Context) (uint64, error)
|
|
ValidateSession(context.Context, string, string) error
|
|
IssueSnapshot(context.Context, string, workerruntime.SnapshotReference) error
|
|
Acknowledge(context.Context, SnapshotAcknowledgement) error
|
|
ReportRuntime(context.Context, workerruntime.Report) (RuntimeDecision, error)
|
|
}
|
|
|
|
func (service *service) CurrentOwnershipEpoch(ctx context.Context) (uint64, error) {
|
|
if ctx == nil {
|
|
return 0, ErrInvalidCommand
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return 0, err
|
|
}
|
|
epoch, err := service.store.CurrentOwnershipEpoch(ctx)
|
|
if err != nil {
|
|
return 0, classifyStoreError(err)
|
|
}
|
|
return epoch, nil
|
|
}
|
|
|
|
func (service *service) ValidateSession(ctx context.Context, workerID, sessionID string) error {
|
|
if ctx == nil || !workerruntime.ValidIdentifier(workerID) || !workerruntime.ValidIdentifier(sessionID) {
|
|
return ErrInvalidCommand
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
if err := service.store.ValidateSession(ctx, workerID, sessionID); err != nil {
|
|
return classifyStoreError(err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (service *service) IssueSnapshot(ctx context.Context, sessionID string, reference workerruntime.SnapshotReference) error {
|
|
if ctx == nil {
|
|
return ErrInvalidCommand
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
normalized, err := workerruntime.NormalizeSnapshotReference(reference)
|
|
if err != nil || !workerruntime.ValidIdentifier(sessionID) {
|
|
return errors.Join(ErrInvalidCommand, err)
|
|
}
|
|
if err := service.store.RecordIssuedSnapshot(ctx, sessionID, normalized, service.options.SessionTTL); err != nil {
|
|
return classifyStoreError(err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
type service struct {
|
|
store workerruntime.ControlStore
|
|
options Options
|
|
snapshots SnapshotSource
|
|
}
|
|
|
|
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
|
|
}
|
|
result := &service{store: store, options: options}
|
|
if options.SnapshotReader != nil {
|
|
source, err := NewOwnedSnapshotSource(
|
|
result, options.SnapshotReader, options.MaxStaleAge, options.MaxRuntimeCounters, options.MaxSnapshotBytes, time.Now,
|
|
)
|
|
if err != nil {
|
|
return nil, errors.Join(ErrInvalidCommand, err)
|
|
}
|
|
result.snapshots = source
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
func (service *service) SnapshotSource() SnapshotSource {
|
|
if service == nil {
|
|
return nil
|
|
}
|
|
return service.snapshots
|
|
}
|
|
|
|
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)
|
|
}
|