93 lines
3.0 KiB
Go
93 lines
3.0 KiB
Go
package pool
|
|
|
|
import (
|
|
"context"
|
|
"time"
|
|
|
|
ownershipDomain "proxy-pool/internal/domain/ownership"
|
|
)
|
|
|
|
var (
|
|
ErrInvalidOwnership = ownershipDomain.ErrInvalidOwnership
|
|
ErrOwnershipUnavailable = ownershipDomain.ErrOwnershipUnavailable
|
|
ErrAlreadyOwned = ownershipDomain.ErrAlreadyOwned
|
|
ErrStaleAssignment = ownershipDomain.ErrStaleAssignment
|
|
ErrNotDraining = ownershipDomain.ErrNotDraining
|
|
ErrDrainNotReady = ownershipDomain.ErrDrainNotReady
|
|
)
|
|
|
|
type Assignment = ownershipDomain.Assignment
|
|
|
|
type OwnershipManager struct {
|
|
repository ownershipDomain.Repository
|
|
}
|
|
|
|
// NewOwnershipManager requires the same authoritative repository used by
|
|
// extraction, so Worker assignment and AVAILABLE -> EXTRACTED cannot race.
|
|
func NewOwnershipManager(repository ownershipDomain.Repository) (*OwnershipManager, error) {
|
|
if repository == nil {
|
|
return nil, ErrInvalidOwnership
|
|
}
|
|
return &OwnershipManager{repository: repository}, nil
|
|
}
|
|
|
|
func (m *OwnershipManager) Assign(ctx context.Context, now time.Time, proxyID, workerID string, ttl time.Duration) (Assignment, error) {
|
|
if m == nil || m.repository == nil || ctx == nil {
|
|
return Assignment{}, ErrInvalidOwnership
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return Assignment{}, err
|
|
}
|
|
return m.repository.Assign(ctx, now, proxyID, workerID, ttl)
|
|
}
|
|
|
|
func (m *OwnershipManager) Renew(ctx context.Context, now time.Time, proxyID, workerID string, epoch uint64, ttl time.Duration) (Assignment, error) {
|
|
if m == nil || m.repository == nil || ctx == nil {
|
|
return Assignment{}, ErrInvalidOwnership
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return Assignment{}, err
|
|
}
|
|
return m.repository.Renew(ctx, now, proxyID, workerID, epoch, ttl)
|
|
}
|
|
|
|
func (m *OwnershipManager) BeginDrain(ctx context.Context, proxyID, workerID string, epoch uint64) (Assignment, error) {
|
|
if m == nil || m.repository == nil || ctx == nil {
|
|
return Assignment{}, ErrInvalidOwnership
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return Assignment{}, err
|
|
}
|
|
return m.repository.BeginDrain(ctx, proxyID, workerID, epoch)
|
|
}
|
|
|
|
func (m *OwnershipManager) AcknowledgeDrain(ctx context.Context, proxyID, workerID string, epoch uint64, active, reserved int64) error {
|
|
if m == nil || m.repository == nil || ctx == nil {
|
|
return ErrInvalidOwnership
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
return m.repository.AcknowledgeDrain(ctx, proxyID, workerID, epoch, active, reserved)
|
|
}
|
|
|
|
func (m *OwnershipManager) Get(ctx context.Context, proxyID string) (Assignment, bool, error) {
|
|
if m == nil || m.repository == nil || ctx == nil {
|
|
return Assignment{}, false, ErrInvalidOwnership
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return Assignment{}, false, err
|
|
}
|
|
return m.repository.Get(ctx, proxyID)
|
|
}
|
|
|
|
func (m *OwnershipManager) Expire(ctx context.Context, now time.Time, limit int) ([]Assignment, error) {
|
|
if m == nil || m.repository == nil || ctx == nil || limit <= 0 {
|
|
return nil, ErrInvalidOwnership
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
return m.repository.Expire(ctx, now, limit)
|
|
}
|