From 80137166399325e39833e519e0ff324f647e174d Mon Sep 17 00:00:00 2001
From: youfak <youfak@126.com>
Date: Sun, 2 Aug 2026 11:39:10 +0800
Subject: [PATCH] fix: rebind drain barriers after worker session change

---
 .../redisactivity/ownership_integration_test.go  |  8 ++++++++
 .../redisactivity/scripts/bind_drain_ticket.lua  | 16 +++++++++-------
 .../domain/activitypool/contracttest/contract.go |  8 ++++++++
 internal/domain/activitypool/pool.go             |  3 ++-
 4 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/internal/adapters/redisactivity/ownership_integration_test.go b/internal/adapters/redisactivity/ownership_integration_test.go
index 1dc8fe5..000414d 100644
--- a/internal/adapters/redisactivity/ownership_integration_test.go
+++ b/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)
diff --git a/internal/adapters/redisactivity/scripts/bind_drain_ticket.lua b/internal/adapters/redisactivity/scripts/bind_drain_ticket.lua
index b0b6906..1540570 100644
--- a/internal/adapters/redisactivity/scripts/bind_drain_ticket.lua
+++ b/internal/adapters/redisactivity/scripts/bind_drain_ticket.lua
@@ -48,14 +48,16 @@ end
 
 local current_epoch = tonumber(ticket.snapshotOwnershipEpoch) or 0
 local current_version = tonumber(ticket.snapshotVersion) or 0
-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
-        return reply('stale')
+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.snapshotChecksum ~= checksum then
+            return reply('stale')
+        end
+        return reply('ok')
     end
-    return reply('ok')
 end
 ticket.sessionId = session_id
 ticket.snapshotVersion = version
diff --git a/internal/domain/activitypool/contracttest/contract.go b/internal/domain/activitypool/contracttest/contract.go
index 9938dd9..f6dcdaf 100644
--- a/internal/domain/activitypool/contracttest/contract.go
+++ b/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)
 	}
diff --git a/internal/domain/activitypool/pool.go b/internal/domain/activitypool/pool.go
index 9a27206..e6c53aa 100644
--- a/internal/domain/activitypool/pool.go
+++ b/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