feat: add redis ownership and expiry maintenance
This commit is contained in:
youfak
parent
57f7c084c9
commit
7aaa1a88c6
91
internal/adapters/redisactivity/maintenance.go
Normal file
91
internal/adapters/redisactivity/maintenance.go
Normal file
@ -0,0 +1,91 @@
|
||||
package redisactivity
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"proxy-pool/internal/domain/activitypool"
|
||||
)
|
||||
|
||||
const (
|
||||
maintenanceInventory = "inventory"
|
||||
maintenanceSweep = "sweep"
|
||||
)
|
||||
|
||||
var (
|
||||
_ activitypool.InventoryReader = (*Adapter)(nil)
|
||||
_ activitypool.Maintainer = (*Adapter)(nil)
|
||||
)
|
||||
|
||||
func (a *Adapter) Inventory(ctx context.Context, upstreamID string, now time.Time) (activitypool.Inventory, error) {
|
||||
result := activitypool.Inventory{UpstreamID: upstreamID}
|
||||
if ctx == nil {
|
||||
return result, activitypool.ErrInvalidInventory
|
||||
}
|
||||
if err := ctx.Err(); err != nil {
|
||||
return result, err
|
||||
}
|
||||
if a == nil || upstreamID == "" || now.IsZero() {
|
||||
return result, activitypool.ErrInvalidInventory
|
||||
}
|
||||
reply, err := a.runMaintenance(ctx, maintenanceInventory, now, a.options.CleanupLimit, upstreamID)
|
||||
if err != nil {
|
||||
return result, err
|
||||
}
|
||||
if reply.Status == scriptInvalid {
|
||||
return result, activitypool.ErrInvalidInventory
|
||||
}
|
||||
if reply.Status != scriptOK || reply.Count < 0 {
|
||||
return result, invalidScriptReply("unexpected inventory reply")
|
||||
}
|
||||
result.Managed = reply.Count
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (a *Adapter) SweepExpired(ctx context.Context, now time.Time, limit int) (int, error) {
|
||||
if ctx == nil {
|
||||
return 0, activitypool.ErrInvalidMaintenance
|
||||
}
|
||||
if err := ctx.Err(); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if a == nil || now.IsZero() || limit <= 0 {
|
||||
return 0, activitypool.ErrInvalidMaintenance
|
||||
}
|
||||
reply, err := a.runMaintenance(ctx, maintenanceSweep, now, limit, "")
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if reply.Status == scriptInvalid {
|
||||
return 0, activitypool.ErrInvalidMaintenance
|
||||
}
|
||||
if reply.Status != scriptOK || reply.Count < 0 || reply.Count > limit {
|
||||
return 0, invalidScriptReply("unexpected expiry sweep reply")
|
||||
}
|
||||
return reply.Count, nil
|
||||
}
|
||||
|
||||
func (a *Adapter) runMaintenance(
|
||||
ctx context.Context,
|
||||
operation string,
|
||||
now time.Time,
|
||||
limit int,
|
||||
upstreamID string,
|
||||
) (maintenanceScriptReply, error) {
|
||||
operationID, err := newOperationID()
|
||||
if err != nil {
|
||||
return maintenanceScriptReply{}, err
|
||||
}
|
||||
result, err := runScript(ctx, a.client, sweepScript, []string{
|
||||
a.keys.records, a.keys.unique, a.keys.idkeys, a.keys.expiry, a.keys.available,
|
||||
a.keys.inventory, a.keys.owners, a.keys.ownerExpiry, a.keys.operation(operationID),
|
||||
}, operation, now.UnixMilli(), limit, upstreamID, operationTTLMillis(a.options.OperationTTL))
|
||||
if err != nil {
|
||||
return maintenanceScriptReply{}, err
|
||||
}
|
||||
var reply maintenanceScriptReply
|
||||
if err := decodeScriptResult(result, &reply); err != nil {
|
||||
return maintenanceScriptReply{}, err
|
||||
}
|
||||
return reply, nil
|
||||
}
|
||||
269
internal/adapters/redisactivity/ownership.go
Normal file
269
internal/adapters/redisactivity/ownership.go
Normal file
@ -0,0 +1,269 @@
|
||||
package redisactivity
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
ownershipDomain "proxy-pool/internal/domain/ownership"
|
||||
)
|
||||
|
||||
const (
|
||||
ownershipAssign = "assign"
|
||||
ownershipRenew = "renew"
|
||||
ownershipBeginDrain = "begin_drain"
|
||||
ownershipAcknowledgeDrain = "acknowledge_drain"
|
||||
ownershipGet = "get"
|
||||
ownershipExpire = "expire"
|
||||
)
|
||||
|
||||
var _ ownershipDomain.Repository = (*Adapter)(nil)
|
||||
|
||||
func (a *Adapter) Assign(
|
||||
ctx context.Context,
|
||||
now time.Time,
|
||||
proxyID string,
|
||||
workerID string,
|
||||
ttl time.Duration,
|
||||
) (ownershipDomain.Assignment, error) {
|
||||
if err := validateOwnershipCall(ctx, a); err != nil {
|
||||
return ownershipDomain.Assignment{}, err
|
||||
}
|
||||
if now.IsZero() || proxyID == "" || workerID == "" || ttl <= 0 {
|
||||
return ownershipDomain.Assignment{}, ownershipDomain.ErrInvalidOwnership
|
||||
}
|
||||
reply, err := a.runOwnership(ctx, ownershipAssign, true, now.UnixMilli(), proxyID, workerID, 0, durationMillis(ttl), 0, 0)
|
||||
if err != nil {
|
||||
return ownershipDomain.Assignment{}, err
|
||||
}
|
||||
switch reply.Status {
|
||||
case scriptOK:
|
||||
return decodeAssignmentReply(reply)
|
||||
case scriptAlreadyOwned:
|
||||
return ownershipDomain.Assignment{}, ownershipDomain.ErrAlreadyOwned
|
||||
case scriptUnavailable, scriptNotFound:
|
||||
return ownershipDomain.Assignment{}, ownershipDomain.ErrOwnershipUnavailable
|
||||
case scriptInvalid:
|
||||
return ownershipDomain.Assignment{}, ownershipDomain.ErrInvalidOwnership
|
||||
default:
|
||||
return ownershipDomain.Assignment{}, invalidScriptReply("unexpected ownership assign status")
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Adapter) Renew(
|
||||
ctx context.Context,
|
||||
now time.Time,
|
||||
proxyID string,
|
||||
workerID string,
|
||||
epoch uint64,
|
||||
ttl time.Duration,
|
||||
) (ownershipDomain.Assignment, error) {
|
||||
if err := validateOwnershipCall(ctx, a); err != nil {
|
||||
return ownershipDomain.Assignment{}, err
|
||||
}
|
||||
if now.IsZero() || proxyID == "" || workerID == "" || epoch == 0 || ttl <= 0 {
|
||||
return ownershipDomain.Assignment{}, ownershipDomain.ErrInvalidOwnership
|
||||
}
|
||||
reply, err := a.runOwnership(ctx, ownershipRenew, true, now.UnixMilli(), proxyID, workerID, epoch, durationMillis(ttl), 0, 0)
|
||||
if err != nil {
|
||||
return ownershipDomain.Assignment{}, err
|
||||
}
|
||||
switch reply.Status {
|
||||
case scriptOK:
|
||||
return decodeAssignmentReply(reply)
|
||||
case scriptStale, scriptNotFound, scriptUnavailable:
|
||||
return ownershipDomain.Assignment{}, ownershipDomain.ErrStaleAssignment
|
||||
case scriptInvalid:
|
||||
return ownershipDomain.Assignment{}, ownershipDomain.ErrInvalidOwnership
|
||||
default:
|
||||
return ownershipDomain.Assignment{}, invalidScriptReply("unexpected ownership renew status")
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Adapter) BeginDrain(
|
||||
ctx context.Context,
|
||||
proxyID string,
|
||||
workerID string,
|
||||
epoch uint64,
|
||||
) (ownershipDomain.Assignment, error) {
|
||||
if err := validateOwnershipCall(ctx, a); err != nil {
|
||||
return ownershipDomain.Assignment{}, err
|
||||
}
|
||||
if proxyID == "" || workerID == "" || epoch == 0 {
|
||||
return ownershipDomain.Assignment{}, ownershipDomain.ErrInvalidOwnership
|
||||
}
|
||||
reply, err := a.runOwnership(ctx, ownershipBeginDrain, true, 0, proxyID, workerID, epoch, 0, 0, 0)
|
||||
if err != nil {
|
||||
return ownershipDomain.Assignment{}, err
|
||||
}
|
||||
switch reply.Status {
|
||||
case scriptOK:
|
||||
return decodeAssignmentReply(reply)
|
||||
case scriptStale, scriptNotFound:
|
||||
return ownershipDomain.Assignment{}, ownershipDomain.ErrStaleAssignment
|
||||
case scriptInvalid:
|
||||
return ownershipDomain.Assignment{}, ownershipDomain.ErrInvalidOwnership
|
||||
default:
|
||||
return ownershipDomain.Assignment{}, invalidScriptReply("unexpected ownership drain status")
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Adapter) AcknowledgeDrain(
|
||||
ctx context.Context,
|
||||
proxyID string,
|
||||
workerID string,
|
||||
epoch uint64,
|
||||
active int64,
|
||||
reserved int64,
|
||||
) error {
|
||||
if err := validateOwnershipCall(ctx, a); err != nil {
|
||||
return err
|
||||
}
|
||||
if proxyID == "" || workerID == "" || epoch == 0 || active < 0 || reserved < 0 {
|
||||
return ownershipDomain.ErrInvalidOwnership
|
||||
}
|
||||
reply, err := a.runOwnership(ctx, ownershipAcknowledgeDrain, true, 0, proxyID, workerID, epoch, 0, active, reserved)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
switch reply.Status {
|
||||
case scriptOK:
|
||||
return nil
|
||||
case scriptStale, scriptNotFound:
|
||||
return ownershipDomain.ErrStaleAssignment
|
||||
case scriptNotDraining:
|
||||
return ownershipDomain.ErrNotDraining
|
||||
case scriptDrainNotReady:
|
||||
return ownershipDomain.ErrDrainNotReady
|
||||
case scriptInvalid:
|
||||
return ownershipDomain.ErrInvalidOwnership
|
||||
default:
|
||||
return invalidScriptReply("unexpected ownership acknowledge status")
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Adapter) Get(ctx context.Context, proxyID string) (ownershipDomain.Assignment, bool, error) {
|
||||
if err := validateOwnershipCall(ctx, a); err != nil {
|
||||
return ownershipDomain.Assignment{}, false, err
|
||||
}
|
||||
if proxyID == "" {
|
||||
return ownershipDomain.Assignment{}, false, ownershipDomain.ErrInvalidOwnership
|
||||
}
|
||||
reply, err := a.runOwnership(ctx, ownershipGet, false, 0, proxyID, "", 0, 0, 0, 0)
|
||||
if err != nil {
|
||||
return ownershipDomain.Assignment{}, false, err
|
||||
}
|
||||
switch reply.Status {
|
||||
case scriptNotFound:
|
||||
return ownershipDomain.Assignment{}, false, nil
|
||||
case scriptOK:
|
||||
assignment, err := decodeAssignmentReply(reply)
|
||||
return assignment, err == nil, err
|
||||
case scriptInvalid:
|
||||
return ownershipDomain.Assignment{}, false, ownershipDomain.ErrInvalidOwnership
|
||||
default:
|
||||
return ownershipDomain.Assignment{}, false, invalidScriptReply("unexpected ownership get status")
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Adapter) Expire(ctx context.Context, now time.Time, limit int) ([]ownershipDomain.Assignment, error) {
|
||||
if err := validateOwnershipCall(ctx, a); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if now.IsZero() || limit <= 0 {
|
||||
return nil, ownershipDomain.ErrInvalidOwnership
|
||||
}
|
||||
reply, err := a.runOwnership(ctx, ownershipExpire, true, now.UnixMilli(), "", "", 0, int64(limit), 0, 0)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if reply.Status == scriptInvalid {
|
||||
return nil, ownershipDomain.ErrInvalidOwnership
|
||||
}
|
||||
if reply.Status != scriptOK || reply.Record == "" {
|
||||
return nil, invalidScriptReply("unexpected ownership expire reply")
|
||||
}
|
||||
var records []ownershipRecord
|
||||
if err := decodeJSON(reply.Record, &records); err != nil {
|
||||
return nil, errors.Join(invalidScriptReply("ownership expire reply contained invalid records"), err)
|
||||
}
|
||||
assignments := make([]ownershipDomain.Assignment, 0, len(records))
|
||||
for _, record := range records {
|
||||
if err := validateOwnershipRecord(record); err != nil {
|
||||
return nil, errors.Join(invalidScriptReply("ownership expire reply contained invalid records"), err)
|
||||
}
|
||||
assignments = append(assignments, assignmentFromRecord(record))
|
||||
}
|
||||
if len(assignments) > limit {
|
||||
return nil, invalidScriptReply("ownership expire reply exceeded limit")
|
||||
}
|
||||
return assignments, nil
|
||||
}
|
||||
|
||||
func (a *Adapter) runOwnership(
|
||||
ctx context.Context,
|
||||
operation string,
|
||||
mutating bool,
|
||||
nowMS int64,
|
||||
proxyID string,
|
||||
workerID string,
|
||||
epoch uint64,
|
||||
value int64,
|
||||
active int64,
|
||||
reserved int64,
|
||||
) (ownershipScriptReply, error) {
|
||||
operationKey := a.keys.epoch
|
||||
if mutating {
|
||||
operationID, err := newOperationID()
|
||||
if err != nil {
|
||||
return ownershipScriptReply{}, err
|
||||
}
|
||||
operationKey = a.keys.operation(operationID)
|
||||
}
|
||||
result, err := runScript(ctx, a.client, ownershipScript, []string{
|
||||
a.keys.records, a.keys.unique, a.keys.idkeys, a.keys.expiry, a.keys.available,
|
||||
a.keys.inventory, a.keys.owners, a.keys.ownerExpiry, a.keys.epoch, operationKey,
|
||||
}, operation, operationTTLMillis(a.options.OperationTTL), a.options.CleanupLimit,
|
||||
nowMS, proxyID, workerID, epoch, value, active, reserved)
|
||||
if err != nil {
|
||||
return ownershipScriptReply{}, err
|
||||
}
|
||||
var reply ownershipScriptReply
|
||||
if err := decodeScriptResult(result, &reply); err != nil {
|
||||
return ownershipScriptReply{}, err
|
||||
}
|
||||
return reply, nil
|
||||
}
|
||||
|
||||
func decodeAssignmentReply(reply ownershipScriptReply) (ownershipDomain.Assignment, error) {
|
||||
if reply.Record == "" {
|
||||
return ownershipDomain.Assignment{}, invalidScriptReply("ownership reply omitted assignment")
|
||||
}
|
||||
record, err := decodeOwnershipRecord(reply.Record)
|
||||
if err != nil {
|
||||
return ownershipDomain.Assignment{}, errors.Join(
|
||||
invalidScriptReply("ownership reply contained an invalid assignment"),
|
||||
fmt.Errorf("decode ownership assignment: %w", err),
|
||||
)
|
||||
}
|
||||
return assignmentFromRecord(record), nil
|
||||
}
|
||||
|
||||
func assignmentFromRecord(record ownershipRecord) ownershipDomain.Assignment {
|
||||
return ownershipDomain.Assignment{
|
||||
ProxyID: record.ProxyID, WorkerID: record.WorkerID, Epoch: record.Epoch,
|
||||
Version: record.AssignmentVersion, ExpiresAt: time.UnixMilli(record.ExpiresAtMS).UTC(),
|
||||
Draining: record.Draining,
|
||||
}
|
||||
}
|
||||
|
||||
func validateOwnershipCall(ctx context.Context, adapter *Adapter) error {
|
||||
if ctx == nil || adapter == nil {
|
||||
return ownershipDomain.ErrInvalidOwnership
|
||||
}
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
408
internal/adapters/redisactivity/ownership_integration_test.go
Normal file
408
internal/adapters/redisactivity/ownership_integration_test.go
Normal file
@ -0,0 +1,408 @@
|
||||
//go:build integration
|
||||
|
||||
package redisactivity
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"proxy-pool/internal/domain/activitypool"
|
||||
extractionDomain "proxy-pool/internal/domain/extraction"
|
||||
ownershipDomain "proxy-pool/internal/domain/ownership"
|
||||
proxyDomain "proxy-pool/internal/domain/proxy"
|
||||
)
|
||||
|
||||
func TestRedisOwnershipLifecycle(t *testing.T) {
|
||||
fixture := newRedisTestFixture(t)
|
||||
now := redisTestNow()
|
||||
seedRedisAvailable(t, fixture.Adapter, "provider-a", now, now.Add(time.Second), 2*time.Minute,
|
||||
testProxy("proxy-a", "192.0.2.10"))
|
||||
|
||||
assigned, err := fixture.Adapter.Assign(context.Background(), now.Add(2*time.Second), "proxy-a", "worker-a", time.Minute)
|
||||
if err != nil || assigned.Epoch == 0 || assigned.Version != 1 || assigned.Draining {
|
||||
t.Fatalf("Assign() = %+v, %v", assigned, err)
|
||||
}
|
||||
assertRedisKeysHaveTTL(t, fixture,
|
||||
fixture.Adapter.keys.owners,
|
||||
fixture.Adapter.keys.ownerExpiry,
|
||||
fixture.Adapter.keys.epoch,
|
||||
)
|
||||
if current, ok, err := fixture.Adapter.Get(context.Background(), "proxy-a"); err != nil || !ok || current != assigned {
|
||||
t.Fatalf("Get(assigned) = %+v, %t, %v", current, ok, err)
|
||||
}
|
||||
blocked, err := fixture.Adapter.Extract(context.Background(), extractionDomain.Command{
|
||||
RequestID: "req-owned", ClientID: "client-a", Requested: 1,
|
||||
Fulfillment: extractionDomain.Partial, Now: now.Add(3 * time.Second),
|
||||
})
|
||||
if err != nil || blocked.Returned != 0 {
|
||||
t.Fatalf("Extract(owned) = %+v, %v", blocked, err)
|
||||
}
|
||||
|
||||
renewed, err := fixture.Adapter.Renew(context.Background(), now.Add(30*time.Second),
|
||||
"proxy-a", "worker-a", assigned.Epoch, 5*time.Minute)
|
||||
if err != nil || renewed.Version != assigned.Version+1 || renewed.Epoch != assigned.Epoch ||
|
||||
!renewed.ExpiresAt.Equal(now.Add(2*time.Minute)) {
|
||||
t.Fatalf("Renew() = %+v, %v", renewed, err)
|
||||
}
|
||||
if _, err := fixture.Adapter.Renew(context.Background(), now.Add(31*time.Second),
|
||||
"proxy-a", "worker-a", assigned.Epoch+1, time.Minute); !errors.Is(err, ownershipDomain.ErrStaleAssignment) {
|
||||
t.Fatalf("Renew(stale epoch) error = %v", err)
|
||||
}
|
||||
|
||||
draining, err := fixture.Adapter.BeginDrain(context.Background(), "proxy-a", "worker-a", assigned.Epoch)
|
||||
if err != nil || !draining.Draining || draining.Version != renewed.Version+1 {
|
||||
t.Fatalf("BeginDrain() = %+v, %v", draining, err)
|
||||
}
|
||||
replayed, err := fixture.Adapter.BeginDrain(context.Background(), "proxy-a", "worker-a", assigned.Epoch)
|
||||
if err != nil || replayed != draining {
|
||||
t.Fatalf("BeginDrain(replay) = %+v, %v", replayed, err)
|
||||
}
|
||||
if err := fixture.Adapter.AcknowledgeDrain(context.Background(), "proxy-a", "worker-a", assigned.Epoch, 1, 0); !errors.Is(err, ownershipDomain.ErrDrainNotReady) {
|
||||
t.Fatalf("AcknowledgeDrain(active) error = %v", err)
|
||||
}
|
||||
if err := fixture.Adapter.AcknowledgeDrain(context.Background(), "proxy-a", "worker-a", assigned.Epoch, 0, 1); !errors.Is(err, ownershipDomain.ErrDrainNotReady) {
|
||||
t.Fatalf("AcknowledgeDrain(reserved) error = %v", err)
|
||||
}
|
||||
if err := fixture.Adapter.AcknowledgeDrain(context.Background(), "proxy-a", "worker-a", assigned.Epoch, 0, 0); err != nil {
|
||||
t.Fatalf("AcknowledgeDrain(): %v", err)
|
||||
}
|
||||
assertRedisKeysHaveTTL(t, fixture,
|
||||
fixture.Adapter.keys.available,
|
||||
fixture.Adapter.keys.protocol("http"),
|
||||
fixture.Adapter.keys.region("cn"),
|
||||
fixture.Adapter.keys.carrier("ct"),
|
||||
fixture.Adapter.keys.upstream("provider-a"),
|
||||
)
|
||||
if current, ok, err := fixture.Adapter.Get(context.Background(), "proxy-a"); err != nil || ok {
|
||||
t.Fatalf("Get(after ACK) = %+v, %t, %v", current, ok, err)
|
||||
}
|
||||
restored, err := fixture.Adapter.Extract(context.Background(), extractionDomain.Command{
|
||||
RequestID: "req-restored", ClientID: "client-a", Requested: 1,
|
||||
Fulfillment: extractionDomain.Partial, Now: now.Add(40 * time.Second),
|
||||
})
|
||||
if err != nil || restored.Returned != 1 || restored.Items[0].ID != "proxy-a" {
|
||||
t.Fatalf("Extract(after ACK) = %+v, %v", restored, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRedisOwnershipRejectsInvalidDrainAndStaleAssignment(t *testing.T) {
|
||||
fixture := newRedisTestFixture(t)
|
||||
now := redisTestNow()
|
||||
seedRedisAvailable(t, fixture.Adapter, "provider-a", now, now.Add(time.Second), time.Minute,
|
||||
testProxy("proxy-a", "192.0.2.10"))
|
||||
assigned, err := fixture.Adapter.Assign(context.Background(), now.Add(2*time.Second), "proxy-a", "worker-a", 20*time.Second)
|
||||
if err != nil {
|
||||
t.Fatalf("Assign(): %v", err)
|
||||
}
|
||||
if _, err := fixture.Adapter.Assign(context.Background(), now.Add(3*time.Second), "proxy-a", "worker-b", time.Minute); !errors.Is(err, ownershipDomain.ErrAlreadyOwned) {
|
||||
t.Fatalf("Assign(already owned) error = %v", err)
|
||||
}
|
||||
if err := fixture.Adapter.AcknowledgeDrain(context.Background(), "proxy-a", "worker-a", assigned.Epoch, 0, 0); !errors.Is(err, ownershipDomain.ErrNotDraining) {
|
||||
t.Fatalf("AcknowledgeDrain(not draining) error = %v", err)
|
||||
}
|
||||
if _, err := fixture.Adapter.BeginDrain(context.Background(), "proxy-a", "worker-b", assigned.Epoch); !errors.Is(err, ownershipDomain.ErrStaleAssignment) {
|
||||
t.Fatalf("BeginDrain(stale worker) error = %v", err)
|
||||
}
|
||||
if err := fixture.Adapter.AcknowledgeDrain(context.Background(), "proxy-a", "worker-a", assigned.Epoch+1, 0, 0); !errors.Is(err, ownershipDomain.ErrStaleAssignment) {
|
||||
t.Fatalf("AcknowledgeDrain(stale epoch) error = %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRedisOwnershipExpireIsLimitedAndAllowsTakeover(t *testing.T) {
|
||||
fixture := newRedisTestFixture(t)
|
||||
now := redisTestNow()
|
||||
for index := range 2 {
|
||||
proxyID := fmt.Sprintf("proxy-%d", index)
|
||||
seedRedisAvailable(t, fixture.Adapter, "provider-a", now, now.Add(time.Second), 2*time.Minute,
|
||||
testProxy(proxyID, fmt.Sprintf("192.0.2.%d", index+10)))
|
||||
if _, err := fixture.Adapter.Assign(context.Background(), now.Add(2*time.Second), proxyID, "worker-a", time.Second); err != nil {
|
||||
t.Fatalf("Assign(%s): %v", proxyID, err)
|
||||
}
|
||||
}
|
||||
first, _, err := fixture.Adapter.Get(context.Background(), "proxy-0")
|
||||
if err != nil {
|
||||
t.Fatalf("Get(proxy-0): %v", err)
|
||||
}
|
||||
expired, err := fixture.Adapter.Expire(context.Background(), now.Add(4*time.Second), 1)
|
||||
if err != nil || len(expired) != 1 {
|
||||
t.Fatalf("Expire(first) = %+v, %v", expired, err)
|
||||
}
|
||||
expired, err = fixture.Adapter.Expire(context.Background(), now.Add(4*time.Second), 1)
|
||||
if err != nil || len(expired) != 1 {
|
||||
t.Fatalf("Expire(second) = %+v, %v", expired, err)
|
||||
}
|
||||
expired, err = fixture.Adapter.Expire(context.Background(), now.Add(4*time.Second), 1)
|
||||
if err != nil || len(expired) != 0 {
|
||||
t.Fatalf("Expire(empty) = %+v, %v", expired, err)
|
||||
}
|
||||
takeover, err := fixture.Adapter.Assign(context.Background(), now.Add(5*time.Second), "proxy-0", "worker-b", time.Minute)
|
||||
if err != nil || takeover.Epoch <= first.Epoch {
|
||||
t.Fatalf("Assign(takeover) = %+v, %v; old epoch=%d", takeover, err, first.Epoch)
|
||||
}
|
||||
|
||||
directFixture := newRedisTestFixture(t)
|
||||
seedRedisAvailable(t, directFixture.Adapter, "provider-a", now, now.Add(time.Second), 2*time.Minute,
|
||||
testProxy("direct", "192.0.2.30"))
|
||||
old, err := directFixture.Adapter.Assign(context.Background(), now.Add(2*time.Second), "direct", "worker-a", time.Second)
|
||||
if err != nil {
|
||||
t.Fatalf("Assign(direct old): %v", err)
|
||||
}
|
||||
direct, err := directFixture.Adapter.Assign(context.Background(), now.Add(4*time.Second), "direct", "worker-b", time.Minute)
|
||||
if err != nil || direct.WorkerID != "worker-b" || direct.Epoch <= old.Epoch {
|
||||
t.Fatalf("Assign(direct takeover) = %+v, %v; old=%+v", direct, err, old)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRedisOwnershipAndExtractionHaveOneWinner(t *testing.T) {
|
||||
fixture := newRedisTestFixture(t)
|
||||
now := redisTestNow()
|
||||
for iteration := range 100 {
|
||||
proxyID := fmt.Sprintf("race-%d", iteration)
|
||||
candidate := testProxy(proxyID, fmt.Sprintf("198.51.100.%d", iteration+1))
|
||||
candidate.State = proxyDomain.StateAvailable
|
||||
if _, err := fixture.Adapter.UpsertFetched(context.Background(), "provider-a", activitypool.FetchedBatch{
|
||||
ObservedAt: now, ConfiguredTTL: 10 * time.Minute, MaxSize: 200,
|
||||
Proxies: []proxyDomain.Proxy{candidate},
|
||||
}); err != nil {
|
||||
t.Fatalf("iteration %d UpsertFetched(): %v", iteration, err)
|
||||
}
|
||||
|
||||
var workers sync.WaitGroup
|
||||
ownershipWon := make(chan bool, 1)
|
||||
extractionWon := make(chan bool, 1)
|
||||
errorsCh := make(chan error, 2)
|
||||
workers.Add(2)
|
||||
go func() {
|
||||
defer workers.Done()
|
||||
_, err := fixture.Adapter.Assign(context.Background(), now.Add(time.Second), proxyID, "worker-a", time.Minute)
|
||||
if err != nil && !errors.Is(err, ownershipDomain.ErrOwnershipUnavailable) {
|
||||
errorsCh <- err
|
||||
}
|
||||
ownershipWon <- err == nil
|
||||
}()
|
||||
go func() {
|
||||
defer workers.Done()
|
||||
result, err := fixture.Adapter.Extract(context.Background(), extractionDomain.Command{
|
||||
RequestID: fmt.Sprintf("req-race-%d", iteration), ClientID: "client-a", Requested: 1,
|
||||
Fulfillment: extractionDomain.Partial, Now: now.Add(time.Second),
|
||||
Upstreams: []string{"provider-a"},
|
||||
})
|
||||
if err != nil {
|
||||
errorsCh <- err
|
||||
}
|
||||
extractionWon <- err == nil && result.Returned == 1 && result.Items[0].ID == proxyID
|
||||
}()
|
||||
workers.Wait()
|
||||
close(errorsCh)
|
||||
for err := range errorsCh {
|
||||
t.Fatalf("iteration %d race error: %v", iteration, err)
|
||||
}
|
||||
winners := 0
|
||||
if <-ownershipWon {
|
||||
winners++
|
||||
}
|
||||
if <-extractionWon {
|
||||
winners++
|
||||
}
|
||||
if winners != 1 {
|
||||
t.Fatalf("iteration %d winners = %d, want 1", iteration, winners)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRedisOwnedProxyIsNotReintroducedByHealth(t *testing.T) {
|
||||
fixture := newRedisTestFixture(t)
|
||||
now := redisTestNow()
|
||||
seedRedisAvailable(t, fixture.Adapter, "provider-a", now, now.Add(time.Second), time.Minute,
|
||||
testProxy("proxy-a", "192.0.2.10"))
|
||||
if _, err := fixture.Adapter.Assign(context.Background(), now.Add(2*time.Second), "proxy-a", "worker-a", 30*time.Second); err != nil {
|
||||
t.Fatalf("Assign(): %v", err)
|
||||
}
|
||||
if _, err := fixture.Adapter.ApplyHealth(context.Background(), activitypool.HealthUpdate{
|
||||
ProxyID: "proxy-a", CheckedAt: now.Add(3 * time.Second), NextState: proxyDomain.StateSuspect,
|
||||
}); err != nil {
|
||||
t.Fatalf("ApplyHealth(suspect): %v", err)
|
||||
}
|
||||
if _, err := fixture.Adapter.ApplyHealth(context.Background(), activitypool.HealthUpdate{
|
||||
ProxyID: "proxy-a", CheckedAt: now.Add(4 * time.Second), NextState: proxyDomain.StateAvailable,
|
||||
}); err != nil {
|
||||
t.Fatalf("ApplyHealth(available): %v", err)
|
||||
}
|
||||
result, err := fixture.Adapter.Extract(context.Background(), extractionDomain.Command{
|
||||
RequestID: "req-owned-health", ClientID: "client-a", Requested: 1,
|
||||
Fulfillment: extractionDomain.Partial, Now: now.Add(5 * time.Second),
|
||||
})
|
||||
if err != nil || result.Returned != 0 {
|
||||
t.Fatalf("Extract(owned after health) = %+v, %v", result, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRedisInventoryTracksExtractionWithoutOwnershipWrites(t *testing.T) {
|
||||
fixture := newRedisTestFixture(t)
|
||||
now := redisTestNow()
|
||||
for index := range 2 {
|
||||
seedRedisAvailable(t, fixture.Adapter, "provider-a", now, now.Add(time.Second), time.Minute,
|
||||
testProxy(fmt.Sprintf("proxy-%d", index), fmt.Sprintf("192.0.2.%d", index+10)))
|
||||
}
|
||||
assertRedisInventory(t, fixture.Adapter, "provider-a", now.Add(2*time.Second), 2)
|
||||
result, err := fixture.Adapter.Extract(context.Background(), extractionDomain.Command{
|
||||
RequestID: "req-inventory", ClientID: "client-a", Requested: 1,
|
||||
Fulfillment: extractionDomain.Partial, Now: now.Add(2 * time.Second),
|
||||
})
|
||||
if err != nil || result.Returned != 1 {
|
||||
t.Fatalf("Extract() = %+v, %v", result, err)
|
||||
}
|
||||
assertRedisInventory(t, fixture.Adapter, "provider-a", now.Add(3*time.Second), 1)
|
||||
remainingID := "proxy-0"
|
||||
if result.Items[0].ID == remainingID {
|
||||
remainingID = "proxy-1"
|
||||
}
|
||||
if _, err := fixture.Adapter.Assign(context.Background(), now.Add(3*time.Second), remainingID, "worker-a", 10*time.Second); err != nil {
|
||||
t.Fatalf("Assign(remaining): %v", err)
|
||||
}
|
||||
assertRedisInventory(t, fixture.Adapter, "provider-a", now.Add(4*time.Second), 1)
|
||||
}
|
||||
|
||||
func TestRedisInventoryPerformsBoundedExpiryCleanup(t *testing.T) {
|
||||
fixture := newRedisTestFixture(t)
|
||||
now := redisTestNow()
|
||||
for _, candidate := range []struct {
|
||||
proxy proxyDomain.Proxy
|
||||
ttl time.Duration
|
||||
}{
|
||||
{proxy: testProxy("short", "192.0.2.10"), ttl: 5 * time.Second},
|
||||
{proxy: testProxy("long", "192.0.2.11"), ttl: time.Minute},
|
||||
} {
|
||||
if _, err := fixture.Adapter.UpsertFetched(context.Background(), "provider-a", activitypool.FetchedBatch{
|
||||
ObservedAt: now, ConfiguredTTL: candidate.ttl, MaxSize: 10,
|
||||
Proxies: []proxyDomain.Proxy{candidate.proxy},
|
||||
}); err != nil {
|
||||
t.Fatalf("UpsertFetched(%s): %v", candidate.proxy.ID, err)
|
||||
}
|
||||
}
|
||||
assertRedisInventory(t, fixture.Adapter, "provider-a", now.Add(6*time.Second), 1)
|
||||
if exists, err := fixture.Client.HExists(context.Background(), fixture.Adapter.keys.records, "short").Result(); err != nil || exists {
|
||||
t.Fatalf("short record after Inventory cleanup = %t, %v", exists, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRedisSweepExpiredIsLimitedAndCleansOwnership(t *testing.T) {
|
||||
fixture := newRedisTestFixture(t)
|
||||
now := redisTestNow()
|
||||
for index := range 2 {
|
||||
candidate := testProxy(fmt.Sprintf("proxy-%d", index), fmt.Sprintf("192.0.2.%d", index+10))
|
||||
if _, err := fixture.Adapter.UpsertFetched(context.Background(), "provider-a", activitypool.FetchedBatch{
|
||||
ObservedAt: now, ConfiguredTTL: 5 * time.Second, MaxSize: 10,
|
||||
Proxies: []proxyDomain.Proxy{candidate},
|
||||
}); err != nil {
|
||||
t.Fatalf("UpsertFetched(proxy-%d): %v", index, err)
|
||||
}
|
||||
}
|
||||
assertRedisInventory(t, fixture.Adapter, "provider-a", now.Add(time.Second), 2)
|
||||
removed, err := fixture.Adapter.SweepExpired(context.Background(), now.Add(6*time.Second), 1)
|
||||
if err != nil || removed != 1 {
|
||||
t.Fatalf("SweepExpired(first) = %d, %v", removed, err)
|
||||
}
|
||||
if remaining, err := fixture.Client.ZCard(context.Background(), fixture.Adapter.keys.expiry).Result(); err != nil || remaining != 1 {
|
||||
t.Fatalf("expiry index after limited sweep = %d, %v; want 1", remaining, err)
|
||||
}
|
||||
removed, err = fixture.Adapter.SweepExpired(context.Background(), now.Add(6*time.Second), 1)
|
||||
if err != nil || removed != 1 {
|
||||
t.Fatalf("SweepExpired(second) = %d, %v", removed, err)
|
||||
}
|
||||
assertRedisInventory(t, fixture.Adapter, "provider-a", now.Add(6*time.Second), 0)
|
||||
removed, err = fixture.Adapter.SweepExpired(context.Background(), now.Add(6*time.Second), 1)
|
||||
if err != nil || removed != 0 {
|
||||
t.Fatalf("SweepExpired(empty) = %d, %v", removed, err)
|
||||
}
|
||||
|
||||
ownedFixture := newRedisTestFixture(t)
|
||||
seedRedisAvailable(t, ownedFixture.Adapter, "provider-a", now, now.Add(time.Second), 5*time.Second,
|
||||
testProxy("owned", "192.0.2.30"))
|
||||
if _, err := ownedFixture.Adapter.Assign(context.Background(), now.Add(2*time.Second), "owned", "worker-a", time.Minute); err != nil {
|
||||
t.Fatalf("Assign(owned): %v", err)
|
||||
}
|
||||
if removed, err := ownedFixture.Adapter.SweepExpired(context.Background(), now.Add(6*time.Second), 1); err != nil || removed != 1 {
|
||||
t.Fatalf("SweepExpired(owned) = %d, %v", removed, err)
|
||||
}
|
||||
if assignment, ok, err := ownedFixture.Adapter.Get(context.Background(), "owned"); err != nil || ok {
|
||||
t.Fatalf("Get(swept owned) = %+v, %t, %v", assignment, ok, err)
|
||||
}
|
||||
assertRedisInventory(t, ownedFixture.Adapter, "provider-a", now.Add(6*time.Second), 0)
|
||||
for _, hashKey := range []string{
|
||||
ownedFixture.Adapter.keys.records,
|
||||
ownedFixture.Adapter.keys.idkeys,
|
||||
ownedFixture.Adapter.keys.owners,
|
||||
} {
|
||||
if exists, err := ownedFixture.Client.HExists(context.Background(), hashKey, "owned").Result(); err != nil || exists {
|
||||
t.Fatalf("hash %s retained owned proxy = %t, %v", hashKey, exists, err)
|
||||
}
|
||||
}
|
||||
if count, err := ownedFixture.Client.HLen(context.Background(), ownedFixture.Adapter.keys.unique).Result(); err != nil || count != 0 {
|
||||
t.Fatalf("unique mappings after sweep = %d, %v", count, err)
|
||||
}
|
||||
for _, sortedSet := range []string{
|
||||
ownedFixture.Adapter.keys.expiry,
|
||||
ownedFixture.Adapter.keys.available,
|
||||
ownedFixture.Adapter.keys.ownerExpiry,
|
||||
ownedFixture.Adapter.keys.protocol("http"),
|
||||
ownedFixture.Adapter.keys.region("cn"),
|
||||
ownedFixture.Adapter.keys.carrier("ct"),
|
||||
ownedFixture.Adapter.keys.upstream("provider-a"),
|
||||
} {
|
||||
if count, err := ownedFixture.Client.ZCard(context.Background(), sortedSet).Result(); err != nil || count != 0 {
|
||||
t.Fatalf("sorted set %s entries after sweep = %d, %v", sortedSet, count, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRedisOwnershipAndMaintenanceValidateInputs(t *testing.T) {
|
||||
fixture := newRedisTestFixture(t)
|
||||
now := redisTestNow()
|
||||
if _, err := fixture.Adapter.Assign(nil, now, "proxy-a", "worker-a", time.Minute); !errors.Is(err, ownershipDomain.ErrInvalidOwnership) {
|
||||
t.Fatalf("Assign(nil context) error = %v", err)
|
||||
}
|
||||
if _, err := fixture.Adapter.Renew(context.Background(), now, "proxy-a", "worker-a", 0, time.Minute); !errors.Is(err, ownershipDomain.ErrInvalidOwnership) {
|
||||
t.Fatalf("Renew(zero epoch) error = %v", err)
|
||||
}
|
||||
if _, err := fixture.Adapter.BeginDrain(context.Background(), "", "worker-a", 1); !errors.Is(err, ownershipDomain.ErrInvalidOwnership) {
|
||||
t.Fatalf("BeginDrain(empty proxy) error = %v", err)
|
||||
}
|
||||
if err := fixture.Adapter.AcknowledgeDrain(context.Background(), "proxy-a", "worker-a", 1, -1, 0); !errors.Is(err, ownershipDomain.ErrInvalidOwnership) {
|
||||
t.Fatalf("AcknowledgeDrain(negative active) error = %v", err)
|
||||
}
|
||||
if _, _, err := fixture.Adapter.Get(context.Background(), ""); !errors.Is(err, ownershipDomain.ErrInvalidOwnership) {
|
||||
t.Fatalf("Get(empty proxy) error = %v", err)
|
||||
}
|
||||
if _, err := fixture.Adapter.Expire(context.Background(), now, 0); !errors.Is(err, ownershipDomain.ErrInvalidOwnership) {
|
||||
t.Fatalf("Expire(zero limit) error = %v", err)
|
||||
}
|
||||
if _, err := fixture.Adapter.Inventory(context.Background(), "", now); !errors.Is(err, activitypool.ErrInvalidInventory) {
|
||||
t.Fatalf("Inventory(empty upstream) error = %v", err)
|
||||
}
|
||||
if _, err := fixture.Adapter.SweepExpired(context.Background(), now, 0); !errors.Is(err, activitypool.ErrInvalidMaintenance) {
|
||||
t.Fatalf("SweepExpired(zero limit) error = %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func assertRedisInventory(t *testing.T, adapter *Adapter, upstreamID string, now time.Time, want int) {
|
||||
t.Helper()
|
||||
inventory, err := adapter.Inventory(context.Background(), upstreamID, now)
|
||||
if err != nil || inventory.UpstreamID != upstreamID || inventory.Managed != want {
|
||||
t.Fatalf("Inventory(%s) = %+v, %v; want %d", upstreamID, inventory, err, want)
|
||||
}
|
||||
}
|
||||
|
||||
func assertRedisKeysHaveTTL(t *testing.T, fixture redisTestFixture, keys ...string) {
|
||||
t.Helper()
|
||||
for _, key := range keys {
|
||||
ttl, err := fixture.Client.PTTL(context.Background(), key).Result()
|
||||
if err != nil || ttl <= 0 {
|
||||
t.Fatalf("PTTL(%s) = %s, %v; want positive TTL", key, ttl, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -17,13 +17,16 @@ import (
|
||||
type scriptStatus string
|
||||
|
||||
const (
|
||||
scriptOK scriptStatus = "ok"
|
||||
scriptInvalid scriptStatus = "invalid"
|
||||
scriptNotFound scriptStatus = "not_found"
|
||||
scriptConflict scriptStatus = "conflict"
|
||||
scriptStale scriptStatus = "stale"
|
||||
scriptUnavailable scriptStatus = "unavailable"
|
||||
scriptInsufficient scriptStatus = "insufficient"
|
||||
scriptOK scriptStatus = "ok"
|
||||
scriptInvalid scriptStatus = "invalid"
|
||||
scriptNotFound scriptStatus = "not_found"
|
||||
scriptConflict scriptStatus = "conflict"
|
||||
scriptStale scriptStatus = "stale"
|
||||
scriptUnavailable scriptStatus = "unavailable"
|
||||
scriptInsufficient scriptStatus = "insufficient"
|
||||
scriptAlreadyOwned scriptStatus = "already_owned"
|
||||
scriptNotDraining scriptStatus = "not_draining"
|
||||
scriptDrainNotReady scriptStatus = "drain_not_ready"
|
||||
)
|
||||
|
||||
type upsertScriptReply struct {
|
||||
@ -64,10 +67,18 @@ var healthSource string
|
||||
//go:embed scripts/extract.lua
|
||||
var extractSource string
|
||||
|
||||
//go:embed scripts/ownership.lua
|
||||
var ownershipSource string
|
||||
|
||||
//go:embed scripts/sweep.lua
|
||||
var sweepSource string
|
||||
|
||||
var (
|
||||
upsertScript = redis.NewScript(upsertSource)
|
||||
healthScript = redis.NewScript(healthSource)
|
||||
extractScript = redis.NewScript(extractSource)
|
||||
upsertScript = redis.NewScript(upsertSource)
|
||||
healthScript = redis.NewScript(healthSource)
|
||||
extractScript = redis.NewScript(extractSource)
|
||||
ownershipScript = redis.NewScript(ownershipSource)
|
||||
sweepScript = redis.NewScript(sweepSource)
|
||||
)
|
||||
|
||||
func runScript(ctx context.Context, client redis.Scripter, script *redis.Script, keys []string, args ...any) (any, error) {
|
||||
|
||||
293
internal/adapters/redisactivity/scripts/ownership.lua
Normal file
293
internal/adapters/redisactivity/scripts/ownership.lua
Normal file
@ -0,0 +1,293 @@
|
||||
local records_key = KEYS[1]
|
||||
local unique_key = KEYS[2]
|
||||
local idkeys_key = KEYS[3]
|
||||
local expiry_key = KEYS[4]
|
||||
local available_key = KEYS[5]
|
||||
local inventory_key = KEYS[6]
|
||||
local owners_key = KEYS[7]
|
||||
local owner_expiry_key = KEYS[8]
|
||||
local epoch_key = KEYS[9]
|
||||
local operation_key = KEYS[10]
|
||||
|
||||
local operation = ARGV[1]
|
||||
local operation_ttl_ms = tonumber(ARGV[2])
|
||||
local cleanup_limit = tonumber(ARGV[3])
|
||||
local now_ms = tonumber(ARGV[4])
|
||||
local proxy_id = ARGV[5]
|
||||
local worker_id = ARGV[6]
|
||||
local epoch = tonumber(ARGV[7])
|
||||
local value = tonumber(ARGV[8])
|
||||
local active = tonumber(ARGV[9])
|
||||
local reserved = tonumber(ARGV[10])
|
||||
local mutating = operation ~= 'get'
|
||||
|
||||
local function finish(reply)
|
||||
local encoded = cjson.encode(reply)
|
||||
if mutating then
|
||||
redis.call('SET', operation_key, encoded, 'PX', operation_ttl_ms)
|
||||
end
|
||||
return encoded
|
||||
end
|
||||
|
||||
if mutating then
|
||||
local committed = redis.call('GET', operation_key)
|
||||
if committed then
|
||||
return committed
|
||||
end
|
||||
end
|
||||
|
||||
local function is_managed(state)
|
||||
return state == 'FETCHED' or state == 'CHECKING' or state == 'AVAILABLE' or
|
||||
state == 'SUSPECT' or state == 'DRAINING'
|
||||
end
|
||||
|
||||
local function decrement_inventory(upstream)
|
||||
if type(upstream) ~= 'string' or upstream == '' then
|
||||
return
|
||||
end
|
||||
local count = redis.call('HINCRBY', inventory_key, upstream, -1)
|
||||
if count < 0 then
|
||||
redis.call('HSET', inventory_key, upstream, 0)
|
||||
end
|
||||
end
|
||||
|
||||
local function touch(key, expires_at_ms)
|
||||
if redis.call('EXISTS', key) == 0 then
|
||||
return
|
||||
end
|
||||
local current = redis.call('PEXPIRETIME', key)
|
||||
if current < expires_at_ms then
|
||||
redis.call('PEXPIREAT', key, expires_at_ms)
|
||||
end
|
||||
end
|
||||
|
||||
local function remove_available(id, record)
|
||||
redis.call('ZREM', available_key, id)
|
||||
local index_keys = record and record.indexKeys
|
||||
if type(index_keys) == 'table' then
|
||||
for _, index_key in ipairs(index_keys) do
|
||||
if type(index_key) == 'string' and index_key ~= '' then
|
||||
redis.call('ZREM', index_key, id)
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
local function add_available(id, record, at_ms)
|
||||
local usable_until_ms = record and tonumber(record.usableUntilMs)
|
||||
if not usable_until_ms or record.state ~= 'AVAILABLE' or usable_until_ms <= at_ms then
|
||||
return
|
||||
end
|
||||
redis.call('ZADD', available_key, usable_until_ms, id)
|
||||
touch(available_key, tonumber(record.expiresAtMs))
|
||||
if type(record.indexKeys) == 'table' then
|
||||
for _, index_key in ipairs(record.indexKeys) do
|
||||
if type(index_key) == 'string' and index_key ~= '' then
|
||||
redis.call('ZADD', index_key, usable_until_ms, id)
|
||||
touch(index_key, tonumber(record.expiresAtMs))
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
local function remove_proxy(id)
|
||||
local raw = redis.call('HGET', records_key, id)
|
||||
local record = nil
|
||||
if raw then
|
||||
local decoded
|
||||
decoded, record = pcall(cjson.decode, raw)
|
||||
remove_available(id, decoded and record or nil)
|
||||
if decoded and type(record) == 'table' and is_managed(record.state) then
|
||||
decrement_inventory(record.sourceUpstream)
|
||||
end
|
||||
else
|
||||
redis.call('ZREM', available_key, id)
|
||||
end
|
||||
local digest = redis.call('HGET', idkeys_key, id)
|
||||
if digest and redis.call('HGET', unique_key, digest) == id then
|
||||
redis.call('HDEL', unique_key, digest)
|
||||
end
|
||||
redis.call('HDEL', idkeys_key, id)
|
||||
redis.call('HDEL', records_key, id)
|
||||
redis.call('ZREM', expiry_key, id)
|
||||
redis.call('HDEL', owners_key, id)
|
||||
redis.call('ZREM', owner_expiry_key, id)
|
||||
end
|
||||
|
||||
local function cleanup_hard_expired(at_ms)
|
||||
local expired = redis.call('ZRANGEBYSCORE', expiry_key, '-inf', at_ms, 'LIMIT', 0, cleanup_limit)
|
||||
for _, id in ipairs(expired) do
|
||||
remove_proxy(id)
|
||||
end
|
||||
end
|
||||
|
||||
local function decode_table(raw)
|
||||
if not raw then
|
||||
return nil
|
||||
end
|
||||
local decoded, value = pcall(cjson.decode, raw)
|
||||
if not decoded or type(value) ~= 'table' then
|
||||
return nil
|
||||
end
|
||||
return value
|
||||
end
|
||||
|
||||
local function valid_assignment(assignment)
|
||||
return assignment and assignment.version == 1 and type(assignment.proxyId) == 'string' and
|
||||
assignment.proxyId ~= '' and type(assignment.workerId) == 'string' and assignment.workerId ~= '' and
|
||||
tonumber(assignment.epoch) and tonumber(assignment.epoch) > 0 and
|
||||
tonumber(assignment.assignmentVersion) and tonumber(assignment.assignmentVersion) > 0 and
|
||||
tonumber(assignment.expiresAtMs) and tonumber(assignment.expiresAtMs) > 0 and
|
||||
type(assignment.draining) == 'boolean'
|
||||
end
|
||||
|
||||
local function clear_owner(id, assignment, at_ms, restore)
|
||||
local raw_record = redis.call('HGET', records_key, id)
|
||||
local record = decode_table(raw_record)
|
||||
if record and (not assignment or record.ownerWorkerId == assignment.workerId) then
|
||||
record.ownerWorkerId = nil
|
||||
redis.call('HSET', records_key, id, cjson.encode(record))
|
||||
if restore then
|
||||
add_available(id, record, at_ms)
|
||||
end
|
||||
end
|
||||
redis.call('HDEL', owners_key, id)
|
||||
redis.call('ZREM', owner_expiry_key, id)
|
||||
end
|
||||
|
||||
if operation == 'assign' then
|
||||
cleanup_hard_expired(now_ms)
|
||||
local current_raw = redis.call('HGET', owners_key, proxy_id)
|
||||
if current_raw then
|
||||
local current = decode_table(current_raw)
|
||||
if not valid_assignment(current) then
|
||||
return finish({status = 'unavailable'})
|
||||
end
|
||||
if tonumber(current.expiresAtMs) > now_ms then
|
||||
return finish({status = 'already_owned'})
|
||||
end
|
||||
clear_owner(proxy_id, current, now_ms, true)
|
||||
end
|
||||
|
||||
local record = decode_table(redis.call('HGET', records_key, proxy_id))
|
||||
if not record or record.state ~= 'AVAILABLE' or
|
||||
(record.ownerWorkerId and record.ownerWorkerId ~= '') or
|
||||
redis.call('HEXISTS', owners_key, proxy_id) == 1 or
|
||||
not tonumber(record.usableUntilMs) or tonumber(record.usableUntilMs) <= now_ms then
|
||||
return finish({status = 'unavailable'})
|
||||
end
|
||||
local expires_at_ms = now_ms + value
|
||||
if tonumber(record.usableUntilMs) < expires_at_ms then
|
||||
expires_at_ms = tonumber(record.usableUntilMs)
|
||||
end
|
||||
local next_epoch = redis.call('INCR', epoch_key)
|
||||
local assignment = {
|
||||
version = 1,
|
||||
proxyId = proxy_id,
|
||||
workerId = worker_id,
|
||||
epoch = next_epoch,
|
||||
assignmentVersion = 1,
|
||||
expiresAtMs = expires_at_ms,
|
||||
draining = false,
|
||||
}
|
||||
local encoded = cjson.encode(assignment)
|
||||
redis.call('HSET', owners_key, proxy_id, encoded)
|
||||
redis.call('ZADD', owner_expiry_key, expires_at_ms, proxy_id)
|
||||
touch(owners_key, tonumber(record.expiresAtMs))
|
||||
touch(owner_expiry_key, tonumber(record.expiresAtMs))
|
||||
touch(epoch_key, tonumber(record.expiresAtMs))
|
||||
record.ownerWorkerId = worker_id
|
||||
redis.call('HSET', records_key, proxy_id, cjson.encode(record))
|
||||
remove_available(proxy_id, record)
|
||||
return finish({status = 'ok', record = encoded})
|
||||
end
|
||||
|
||||
if operation == 'renew' then
|
||||
cleanup_hard_expired(now_ms)
|
||||
local current = decode_table(redis.call('HGET', owners_key, proxy_id))
|
||||
if not valid_assignment(current) or current.workerId ~= worker_id or tonumber(current.epoch) ~= epoch then
|
||||
return finish({status = 'stale'})
|
||||
end
|
||||
if tonumber(current.expiresAtMs) <= now_ms then
|
||||
clear_owner(proxy_id, current, now_ms, true)
|
||||
return finish({status = 'stale'})
|
||||
end
|
||||
local record = decode_table(redis.call('HGET', records_key, proxy_id))
|
||||
if not record or record.ownerWorkerId ~= worker_id or
|
||||
not tonumber(record.usableUntilMs) or tonumber(record.usableUntilMs) <= now_ms then
|
||||
clear_owner(proxy_id, current, now_ms, false)
|
||||
return finish({status = 'stale'})
|
||||
end
|
||||
local expires_at_ms = now_ms + value
|
||||
if tonumber(record.usableUntilMs) < expires_at_ms then
|
||||
expires_at_ms = tonumber(record.usableUntilMs)
|
||||
end
|
||||
current.assignmentVersion = tonumber(current.assignmentVersion) + 1
|
||||
current.expiresAtMs = expires_at_ms
|
||||
local encoded = cjson.encode(current)
|
||||
redis.call('HSET', owners_key, proxy_id, encoded)
|
||||
redis.call('ZADD', owner_expiry_key, expires_at_ms, proxy_id)
|
||||
touch(owners_key, tonumber(record.expiresAtMs))
|
||||
touch(owner_expiry_key, tonumber(record.expiresAtMs))
|
||||
touch(epoch_key, tonumber(record.expiresAtMs))
|
||||
return finish({status = 'ok', record = encoded})
|
||||
end
|
||||
|
||||
if operation == 'begin_drain' then
|
||||
local current = decode_table(redis.call('HGET', owners_key, proxy_id))
|
||||
if not valid_assignment(current) or current.workerId ~= worker_id or tonumber(current.epoch) ~= epoch then
|
||||
return finish({status = 'stale'})
|
||||
end
|
||||
if not current.draining then
|
||||
current.draining = true
|
||||
current.assignmentVersion = tonumber(current.assignmentVersion) + 1
|
||||
local encoded = cjson.encode(current)
|
||||
redis.call('HSET', owners_key, proxy_id, encoded)
|
||||
return finish({status = 'ok', record = encoded})
|
||||
end
|
||||
return finish({status = 'ok', record = cjson.encode(current)})
|
||||
end
|
||||
|
||||
if operation == 'acknowledge_drain' then
|
||||
local current = decode_table(redis.call('HGET', owners_key, proxy_id))
|
||||
if not valid_assignment(current) or current.workerId ~= worker_id or tonumber(current.epoch) ~= epoch then
|
||||
return finish({status = 'stale'})
|
||||
end
|
||||
if not current.draining then
|
||||
return finish({status = 'not_draining'})
|
||||
end
|
||||
if active > 0 or reserved > 0 then
|
||||
return finish({status = 'drain_not_ready'})
|
||||
end
|
||||
local server_time = redis.call('TIME')
|
||||
local server_now_ms = tonumber(server_time[1]) * 1000 + math.floor(tonumber(server_time[2]) / 1000)
|
||||
clear_owner(proxy_id, current, server_now_ms, true)
|
||||
return finish({status = 'ok'})
|
||||
end
|
||||
|
||||
if operation == 'get' then
|
||||
local raw = redis.call('HGET', owners_key, proxy_id)
|
||||
if not raw then
|
||||
return finish({status = 'not_found'})
|
||||
end
|
||||
return finish({status = 'ok', record = raw})
|
||||
end
|
||||
|
||||
if operation == 'expire' then
|
||||
local ids = redis.call('ZRANGEBYSCORE', owner_expiry_key, '-inf', now_ms, 'LIMIT', 0, value)
|
||||
local expired = {}
|
||||
for _, id in ipairs(ids) do
|
||||
local current = decode_table(redis.call('HGET', owners_key, id))
|
||||
if valid_assignment(current) then
|
||||
expired[#expired + 1] = current
|
||||
end
|
||||
clear_owner(id, current, now_ms, true)
|
||||
end
|
||||
local encoded = '[]'
|
||||
if #expired > 0 then
|
||||
encoded = cjson.encode(expired)
|
||||
end
|
||||
return finish({status = 'ok', record = encoded})
|
||||
end
|
||||
|
||||
return finish({status = 'invalid'})
|
||||
95
internal/adapters/redisactivity/scripts/sweep.lua
Normal file
95
internal/adapters/redisactivity/scripts/sweep.lua
Normal file
@ -0,0 +1,95 @@
|
||||
local records_key = KEYS[1]
|
||||
local unique_key = KEYS[2]
|
||||
local idkeys_key = KEYS[3]
|
||||
local expiry_key = KEYS[4]
|
||||
local available_key = KEYS[5]
|
||||
local inventory_key = KEYS[6]
|
||||
local owners_key = KEYS[7]
|
||||
local owner_expiry_key = KEYS[8]
|
||||
local operation_key = KEYS[9]
|
||||
|
||||
local operation = ARGV[1]
|
||||
local now_ms = tonumber(ARGV[2])
|
||||
local limit = tonumber(ARGV[3])
|
||||
local upstream_id = ARGV[4]
|
||||
local operation_ttl_ms = tonumber(ARGV[5])
|
||||
|
||||
local function finish(reply)
|
||||
local encoded = cjson.encode(reply)
|
||||
redis.call('SET', operation_key, encoded, 'PX', operation_ttl_ms)
|
||||
return encoded
|
||||
end
|
||||
|
||||
local committed = redis.call('GET', operation_key)
|
||||
if committed then
|
||||
return committed
|
||||
end
|
||||
|
||||
local function is_managed(state)
|
||||
return state == 'FETCHED' or state == 'CHECKING' or state == 'AVAILABLE' or
|
||||
state == 'SUSPECT' or state == 'DRAINING'
|
||||
end
|
||||
|
||||
local function decrement_inventory(upstream)
|
||||
if type(upstream) ~= 'string' or upstream == '' then
|
||||
return
|
||||
end
|
||||
local count = redis.call('HINCRBY', inventory_key, upstream, -1)
|
||||
if count < 0 then
|
||||
redis.call('HSET', inventory_key, upstream, 0)
|
||||
end
|
||||
end
|
||||
|
||||
local function remove_available(proxy_id, record)
|
||||
redis.call('ZREM', available_key, proxy_id)
|
||||
local index_keys = record and record.indexKeys
|
||||
if type(index_keys) == 'table' then
|
||||
for _, index_key in ipairs(index_keys) do
|
||||
if type(index_key) == 'string' and index_key ~= '' then
|
||||
redis.call('ZREM', index_key, proxy_id)
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
local function remove_proxy(proxy_id)
|
||||
local raw = redis.call('HGET', records_key, proxy_id)
|
||||
local record = nil
|
||||
if raw then
|
||||
local decoded
|
||||
decoded, record = pcall(cjson.decode, raw)
|
||||
remove_available(proxy_id, decoded and record or nil)
|
||||
if decoded and type(record) == 'table' and is_managed(record.state) then
|
||||
decrement_inventory(record.sourceUpstream)
|
||||
end
|
||||
else
|
||||
redis.call('ZREM', available_key, proxy_id)
|
||||
end
|
||||
local digest = redis.call('HGET', idkeys_key, proxy_id)
|
||||
if digest and redis.call('HGET', unique_key, digest) == proxy_id then
|
||||
redis.call('HDEL', unique_key, digest)
|
||||
end
|
||||
redis.call('HDEL', idkeys_key, proxy_id)
|
||||
redis.call('HDEL', records_key, proxy_id)
|
||||
redis.call('ZREM', expiry_key, proxy_id)
|
||||
redis.call('HDEL', owners_key, proxy_id)
|
||||
redis.call('ZREM', owner_expiry_key, proxy_id)
|
||||
end
|
||||
|
||||
local expired = redis.call('ZRANGEBYSCORE', expiry_key, '-inf', now_ms, 'LIMIT', 0, limit)
|
||||
for _, proxy_id in ipairs(expired) do
|
||||
remove_proxy(proxy_id)
|
||||
end
|
||||
|
||||
if operation == 'sweep' then
|
||||
return finish({status = 'ok', count = #expired})
|
||||
end
|
||||
if operation == 'inventory' then
|
||||
local count = tonumber(redis.call('HGET', inventory_key, upstream_id) or '0')
|
||||
if count < 0 then
|
||||
count = 0
|
||||
redis.call('HSET', inventory_key, upstream_id, 0)
|
||||
end
|
||||
return finish({status = 'ok', count = count})
|
||||
end
|
||||
return finish({status = 'invalid', count = 0})
|
||||
Loading…
Reference in New Issue
Block a user