From ab87154a2eef52fa7cd50eae692bbeb224a2a79d Mon Sep 17 00:00:00 2001
From: Pallav Garg <pallav.garg+css@cloudkitchens.com>
Date: Fri, 31 Jul 2026 13:56:30 -0700
Subject: [PATCH] Prevent leader from acknowledging failed writes

---
 pkg/service/leader/leader_test.go | 33 +++++++++++++++++++++++++++----
 pkg/service/leader/writer.go      | 17 ++++++++--------
 2 files changed, 37 insertions(+), 13 deletions(-)

diff --git a/pkg/service/leader/leader_test.go b/pkg/service/leader/leader_test.go
index 631afad1..7fe2af01 100644
--- a/pkg/service/leader/leader_test.go
+++ b/pkg/service/leader/leader_test.go
@@ -2,6 +2,7 @@ package leader_test
 
 import (
 	"context"
+	"errors"
 	"testing"
 	"testing/synctest"
 	"time"
@@ -9,17 +10,17 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"go.atoms.co/lib/chanx"
+	"go.atoms.co/lib/testing/assertx"
 	"go.atoms.co/splitter/lib/service/location"
 	"go.atoms.co/splitter/lib/service/session"
-	"go.atoms.co/lib/testing/assertx"
+	splitterpb "go.atoms.co/splitter/pb"
+	splitterprivatepb "go.atoms.co/splitter/pb/private"
 	"go.atoms.co/splitter/pkg/core"
 	"go.atoms.co/splitter/pkg/model"
 	"go.atoms.co/splitter/pkg/service/leader"
 	"go.atoms.co/splitter/pkg/storage"
 	"go.atoms.co/splitter/pkg/storage/memory"
-	splitterprivatepb "go.atoms.co/splitter/pb/private"
-	splitterpb "go.atoms.co/splitter/pb"
-	"go.atoms.co/lib/chanx"
 )
 
 const (
@@ -190,6 +191,22 @@ func TestLeader_Operations(t *testing.T) {
 	assert.Len(t, snap.GetSnapshot().GetTenants(), 2)
 }
 
+func TestLeader_DoesNotAcknowledgeFailedUpdate(t *testing.T) {
+	ctx := context.Background()
+	loc := location.New("centralus", "splitter-0")
+	db := failingUpdateStorage{Storage: memory.New()}
+
+	l := leader.New(ctx, loc, db, 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.ErrorIs(t, err, model.ErrNotOwned)
+	require.Nil(t, response)
+}
+
 func TestLeader_HandleUpdate(t *testing.T) {
 	synctest.Test(t, func(t *testing.T) {
 		ctx := context.Background()
@@ -283,6 +300,14 @@ func setup(t *testing.T, ctx context.Context, services ...model.Service) storage
 	return db
 }
 
+type failingUpdateStorage struct {
+	storage.Storage
+}
+
+func (failingUpdateStorage) Update(context.Context, core.Update) error {
+	return errors.New("apply failed")
+}
+
 func setupWithDomains(t *testing.T, ctx context.Context, service model.Service, domains ...model.Domain) storage.Storage {
 	db := memory.New()
 
diff --git a/pkg/service/leader/writer.go b/pkg/service/leader/writer.go
index 17d0ce2e..f8ae8ae0 100644
--- a/pkg/service/leader/writer.go
+++ b/pkg/service/leader/writer.go
@@ -5,14 +5,14 @@ import (
 	"fmt"
 	"time"
 
-	"go.atoms.co/lib/log"
 	"go.atoms.co/iox"
+	"go.atoms.co/lib/log"
 	"go.atoms.co/lib/workqueue"
+	splitterpb "go.atoms.co/splitter/pb"
+	splitterprivatepb "go.atoms.co/splitter/pb/private"
 	"go.atoms.co/splitter/pkg/core"
 	"go.atoms.co/splitter/pkg/model"
 	"go.atoms.co/splitter/pkg/storage"
-	splitterprivatepb "go.atoms.co/splitter/pb/private"
-	splitterpb "go.atoms.co/splitter/pb"
 )
 
 const (
@@ -612,8 +612,6 @@ func (w *Writer) updateAsync(ctx context.Context, upd core.Update) iox.AsyncClos
 func (w *Writer) applyUpdateAsync(ctx context.Context, upd core.Update) iox.AsyncCloser {
 	done := iox.NewAsyncCloser()
 	w.pool.Chan() <- func() {
-		defer done.Close()
-
 		// Perform I/O async. If it fails, escalate.
 
 		if err := w.db.Update(ctx, upd); err != nil {
@@ -627,6 +625,7 @@ func (w *Writer) applyUpdateAsync(ctx context.Context, upd core.Update) iox.Asyn
 		case <-w.Closed():
 		}
 
+		done.Close()
 	}
 	return done
 }
@@ -654,8 +653,6 @@ func (w *Writer) deleteAsync(ctx context.Context, del core.Delete) iox.AsyncClos
 
 	done := iox.NewAsyncCloser()
 	w.pool.Chan() <- func() {
-		defer done.Close()
-
 		// Perform I/O async. If it fails, escalate.
 
 		if err := w.db.Delete(ctx, del); err != nil {
@@ -668,6 +665,8 @@ func (w *Writer) deleteAsync(ctx context.Context, del core.Delete) iox.AsyncClos
 		case w.del <- del:
 		case <-w.Closed():
 		}
+
+		done.Close()
 	}
 	return done
 }
@@ -677,8 +676,6 @@ func (w *Writer) restoreAsync(ctx context.Context, res core.Restore) iox.AsyncCl
 
 	done := iox.NewAsyncCloser()
 	w.pool.Chan() <- func() {
-		defer done.Close()
-
 		// Perform I/O async. If it fails, escalate.
 
 		if err := w.db.Restore(ctx, res); err != nil {
@@ -691,6 +688,8 @@ func (w *Writer) restoreAsync(ctx context.Context, res core.Restore) iox.AsyncCl
 		case w.res <- res:
 		case <-w.Closed():
 		}
+
+		done.Close()
 	}
 	return done
 }