From 7aaa1a88c6aeaafb8c0a8e655d44f90689e80385 Mon Sep 17 00:00:00 2001 From: youfak Date: Wed, 29 Jul 2026 17:33:40 +0800 Subject: [PATCH] feat: add redis ownership and expiry maintenance --- .../adapters/redisactivity/maintenance.go | 91 ++++ internal/adapters/redisactivity/ownership.go | 269 ++++++++++++ .../ownership_integration_test.go | 408 ++++++++++++++++++ internal/adapters/redisactivity/scripts.go | 31 +- .../redisactivity/scripts/ownership.lua | 293 +++++++++++++ .../adapters/redisactivity/scripts/sweep.lua | 95 ++++ 6 files changed, 1177 insertions(+), 10 deletions(-) create mode 100644 internal/adapters/redisactivity/maintenance.go create mode 100644 internal/adapters/redisactivity/ownership.go create mode 100644 internal/adapters/redisactivity/ownership_integration_test.go create mode 100644 internal/adapters/redisactivity/scripts/ownership.lua create mode 100644 internal/adapters/redisactivity/scripts/sweep.lua diff --git a/internal/adapters/redisactivity/maintenance.go b/internal/adapters/redisactivity/maintenance.go new file mode 100644 index 0000000..3aa6572 --- /dev/null +++ b/internal/adapters/redisactivity/maintenance.go @@ -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 +} diff --git a/internal/adapters/redisactivity/ownership.go b/internal/adapters/redisactivity/ownership.go new file mode 100644 index 0000000..9af9dfe --- /dev/null +++ b/internal/adapters/redisactivity/ownership.go @@ -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 +} diff --git a/internal/adapters/redisactivity/ownership_integration_test.go b/internal/adapters/redisactivity/ownership_integration_test.go new file mode 100644 index 0000000..e2881c7 --- /dev/null +++ b/internal/adapters/redisactivity/ownership_integration_test.go @@ -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) + } + } +} diff --git a/internal/adapters/redisactivity/scripts.go b/internal/adapters/redisactivity/scripts.go index 1bf04ac..7067498 100644 --- a/internal/adapters/redisactivity/scripts.go +++ b/internal/adapters/redisactivity/scripts.go @@ -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) { diff --git a/internal/adapters/redisactivity/scripts/ownership.lua b/internal/adapters/redisactivity/scripts/ownership.lua new file mode 100644 index 0000000..3506591 --- /dev/null +++ b/internal/adapters/redisactivity/scripts/ownership.lua @@ -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'}) diff --git a/internal/adapters/redisactivity/scripts/sweep.lua b/internal/adapters/redisactivity/scripts/sweep.lua new file mode 100644 index 0000000..342f8e5 --- /dev/null +++ b/internal/adapters/redisactivity/scripts/sweep.lua @@ -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})