From 23b66643a0fa9457129c09e07a98abf4167b0da3 Mon Sep 17 00:00:00 2001
From: Pallav Garg <pallav.garg+css@cloudkitchens.com>
Date: Fri, 31 Jul 2026 14:08:12 -0700
Subject: [PATCH] Reset writer state after restore

---
 pkg/service/leader/leader_test.go | 27 +++++++++++++++++++++++++++
 pkg/service/leader/writer.go      |  4 +++-
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/pkg/service/leader/leader_test.go b/pkg/service/leader/leader_test.go
index 7eacc05..dec68a1 100644
--- a/pkg/service/leader/leader_test.go
+++ b/pkg/service/leader/leader_test.go
@@ -191,6 +191,33 @@ func TestLeader_Operations(t *testing.T) {
 	assert.Len(t, snap.GetSnapshot().GetTenants(), 2)
 }
 
+func TestLeader_RestoreResetsWriterState(t *testing.T) {
+	ctx := context.Background()
+	loc := location.New("centralus", "splitter-0")
+
+	l := leader.New(ctx, loc, memory.New(), leader.WithFastActivation())
+	defer l.Close()
+	<-l.Initialized().Closed()
+
+	response, err := l.Handle(ctx, leader.NewHandleTenantRequest(&splitterprivatepb.TenantRequest{
+		Req: &splitterprivatepb.TenantRequest_New{New: &splitterpb.NewTenantRequest{Name: string(tenant1)}},
+	}))
+	require.NoError(t, err)
+	require.NotNil(t, response.GetTenant().GetNew())
+
+	response, err = l.Handle(ctx, leader.NewHandleOperationRequest(&splitterprivatepb.OperationRequest{
+		Req: &splitterprivatepb.OperationRequest_Restore{Restore: &splitterprivatepb.RestoreRequest{Nuke: true}},
+	}))
+	require.NoError(t, err)
+	require.NotNil(t, response.GetOperation().GetRestore())
+
+	response, err = l.Handle(ctx, leader.NewHandleTenantRequest(&splitterprivatepb.TenantRequest{
+		Req: &splitterprivatepb.TenantRequest_New{New: &splitterpb.NewTenantRequest{Name: string(tenant1)}},
+	}))
+	require.NoError(t, err)
+	require.NotNil(t, response.GetTenant().GetNew())
+}
+
 func TestLeader_DoesNotAcknowledgeFailedUpdate(t *testing.T) {
 	ctx := context.Background()
 	loc := location.New("centralus", "splitter-0")
diff --git a/pkg/service/leader/writer.go b/pkg/service/leader/writer.go
index 9a49770..0e74bed 100644
--- a/pkg/service/leader/writer.go
+++ b/pkg/service/leader/writer.go
@@ -672,7 +672,9 @@ func (w *Writer) deleteAsync(ctx context.Context, del core.Delete) iox.AsyncClos
 }
 
 func (w *Writer) restoreAsync(ctx context.Context, res core.Restore) iox.AsyncCloser {
-	w.cache.Restore(res.Snapshot())
+	snapshot := res.Snapshot()
+	w.writer = storage.NewWriter(snapshot)
+	w.cache.Restore(snapshot)
 
 	done := iox.NewAsyncCloser()
 	w.pool.Chan() <- func() {