76 lines
2.2 KiB
Go
76 lines
2.2 KiB
Go
package worker
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"time"
|
|
|
|
controlplanev1 "proxy-pool/gen/controlplane/v1"
|
|
"proxy-pool/internal/controlplane/snapshotwire"
|
|
"proxy-pool/internal/domain/workerruntime"
|
|
|
|
"google.golang.org/protobuf/types/known/timestamppb"
|
|
)
|
|
|
|
var ErrSnapshotsUnavailable = errors.New("worker snapshots are unavailable")
|
|
|
|
type SnapshotWatchRequest struct {
|
|
WorkerID string
|
|
SessionID string
|
|
LastAppliedVersion uint64
|
|
LastChecksum []byte
|
|
}
|
|
|
|
type SnapshotSource interface {
|
|
Watch(context.Context, SnapshotWatchRequest) (<-chan *controlplanev1.WorkerSnapshot, error)
|
|
}
|
|
|
|
type OwnershipEpochReader interface {
|
|
CurrentOwnershipEpoch(context.Context) (uint64, error)
|
|
}
|
|
|
|
type InitialSnapshotSource struct {
|
|
epochs OwnershipEpochReader
|
|
validFor time.Duration
|
|
now func() time.Time
|
|
}
|
|
|
|
func NewInitialSnapshotSource(epochs OwnershipEpochReader, validFor time.Duration, now func() time.Time) (*InitialSnapshotSource, error) {
|
|
if epochs == nil || validFor <= 0 || now == nil {
|
|
return nil, ErrSnapshotsUnavailable
|
|
}
|
|
return &InitialSnapshotSource{epochs: epochs, validFor: validFor, now: now}, nil
|
|
}
|
|
|
|
func (source *InitialSnapshotSource) Watch(ctx context.Context, request SnapshotWatchRequest) (<-chan *controlplanev1.WorkerSnapshot, error) {
|
|
if source == nil || ctx == nil || !workerruntime.ValidIdentifier(request.WorkerID) ||
|
|
!workerruntime.ValidIdentifier(request.SessionID) || request.LastAppliedVersion == ^uint64(0) {
|
|
return nil, ErrSnapshotsUnavailable
|
|
}
|
|
epoch, err := source.epochs.CurrentOwnershipEpoch(ctx)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if epoch == 0 {
|
|
return nil, ErrSnapshotsUnavailable
|
|
}
|
|
now := source.now().UTC()
|
|
if now.IsZero() {
|
|
return nil, ErrSnapshotsUnavailable
|
|
}
|
|
full := &controlplanev1.WorkerSnapshot{
|
|
Version: request.LastAppliedVersion + 1, OwnershipEpoch: epoch,
|
|
GeneratedAt: timestamppb.New(now), ValidUntil: timestamppb.New(now.Add(source.validFor)),
|
|
}
|
|
checksum, err := snapshotwire.Checksum(full)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
full.Checksum = append([]byte(nil), checksum[:]...)
|
|
updates := make(chan *controlplanev1.WorkerSnapshot, 1)
|
|
updates <- full
|
|
return updates, nil
|
|
}
|
|
|
|
var _ SnapshotSource = (*InitialSnapshotSource)(nil)
|