fix: rebind drain barriers after worker session change

internal/adapters/redisactivity/ownership_integration_test.go

@@ -82,6 +82,14 @@ func TestRedisOwnershipLifecycle(t *testing.T) {
 	if bound, err := fixture.Adapter.PendingDrains(context.Background(), "worker-a", 1); err != nil || len(bound) != 1 || bound[0].Barrier != pending[0].Barrier {
 		t.Fatalf("PendingDrains(bound) = %+v, %v", bound, err)
 	}
+	pending[0].Barrier.SessionID = "session-b"
+	pending[0].Barrier.Version = 1
+	if err := fixture.Adapter.BindDrainBarrier(context.Background(), pending[0]); err != nil {
+		t.Fatalf("BindDrainBarrier(new session): %v", err)
+	}
+	if rebound, err := fixture.Adapter.PendingDrains(context.Background(), "worker-a", 1); err != nil || len(rebound) != 1 || rebound[0].Barrier != pending[0].Barrier {
+		t.Fatalf("PendingDrains(rebound) = %+v, %v", rebound, 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)

internal/adapters/redisactivity/scripts/bind_drain_ticket.lua

@@ -48,15 +48,17 @@ end
 
 local current_epoch = tonumber(ticket.snapshotOwnershipEpoch) or 0
 local current_version = tonumber(ticket.snapshotVersion) or 0
+if ticket.sessionId == session_id then
     if current_epoch > snapshot_epoch or (current_epoch == snapshot_epoch and current_version > version) then
         return reply('ok')
     end
     if current_epoch == snapshot_epoch and current_version == version then
-    if ticket.sessionId ~= session_id or ticket.snapshotChecksum ~= checksum then
+        if ticket.snapshotChecksum ~= checksum then
             return reply('stale')
         end
         return reply('ok')
     end
+end
 ticket.sessionId = session_id
 ticket.snapshotVersion = version
 ticket.snapshotOwnershipEpoch = snapshot_epoch

internal/domain/activitypool/contracttest/contract.go

@@ -310,6 +310,14 @@ func runOwnershipContract(t *testing.T, factory Factory) {
 	if err != nil || len(bound) != 1 || bound[0].Barrier != pending[0].Barrier {
 		t.Fatalf("PendingDrains(bound) = %+v, %v", bound, err)
 	}
+	pending[0].Barrier.SessionID = "session-b"
+	pending[0].Barrier.Version = 1
+	if err := store.BindDrainBarrier(context.Background(), pending[0]); err != nil {
+		t.Fatalf("BindDrainBarrier(new session): %v", err)
+	}
+	if rebound, err := store.PendingDrains(context.Background(), "worker-a", 1); err != nil || len(rebound) != 1 || rebound[0].Barrier != pending[0].Barrier {
+		t.Fatalf("PendingDrains(rebound) = %+v, %v", rebound, err)
+	}
 	if err := store.AcknowledgeDrain(context.Background(), "proxy-a", "worker-a", assigned.Epoch, 1, 0); !errors.Is(err, ownershipDomain.ErrDrainNotReady) {
 		t.Fatalf("AcknowledgeDrain(active) error = %v", err)
 	}

internal/domain/activitypool/pool.go

@@ -950,7 +950,8 @@ func (p *MemoryPool) BindDrainBarrier(ctx context.Context, ticket ownershipDomai
 		current.RequiredSnapshotEpoch != ticket.RequiredSnapshotEpoch {
 		return ownershipDomain.ErrStaleAssignment
 	}
-	if barrierAfter(current.Barrier, ticket.Barrier) || barrierEqual(current.Barrier, ticket.Barrier) {
+	if current.Barrier.SessionID == ticket.Barrier.SessionID &&
+		(barrierAfter(current.Barrier, ticket.Barrier) || barrierEqual(current.Barrier, ticket.Barrier)) {
 		return nil
 	}
 	current.Barrier = ticket.Barrier