fix(disttae): retain unpublished S3 cleanup ownership

pkg/cnservice/server.go

@@ -544,6 +544,8 @@ func (s *service) closeService() error {
 			// otherwise retirement can wait for the work we have not stopped yet.
 			s.closeSiriusRuntime,
 			s.closeMongoDBRuntime,
+			// All transaction producers must stop before cleanup admission closes.
+			func() error { return s.colexecServer.CloseUnpublishedS3Cleanup(context.Background()) },
 		)
 		if s.closeErr != nil {
 			return
@@ -860,14 +862,7 @@ func (s *service) handleRequest(
 		}
 	}
 
-	// start a goroutine to handle one received message.
-	owned = false
-	cancelOwned = false
-	go func() {
-		defer release()
-		if value.Cancel != nil {
-			defer value.Cancel()
-		}
+	invoke := func() {
 		s.pipelines.counter.Add(1)
 		defer s.pipelines.counter.Add(-1)
 
@@ -885,6 +880,24 @@ func (s *service) handleRequest(
 			s._txnClient,
 			s.aicm,
 			s.acquireMessage)
+	}
+	// The connection read loop calls handleRequest in wire order. Ownership
+	// receipts are per batch, so ACKs must not be processed out of order by
+	// separate handler goroutines.
+	if msg.GetCmd() == pipeline.Method_PipelineBatchAck {
+		invoke()
+		return nil
+	}
+
+	// start a goroutine to handle one received message.
+	owned = false
+	cancelOwned = false
+	go func() {
+		defer release()
+		if value.Cancel != nil {
+			defer value.Cancel()
+		}
+		invoke()
 	}()
 	return nil
 }

pkg/cnservice/server_test.go

@@ -56,6 +56,7 @@ import (
 	"github.com/matrixorigin/matrixone/pkg/pb/timestamp"
 	"github.com/matrixorigin/matrixone/pkg/queryservice"
 	qclient "github.com/matrixorigin/matrixone/pkg/queryservice/client"
+	"github.com/matrixorigin/matrixone/pkg/sql/colexec"
 	"github.com/matrixorigin/matrixone/pkg/sql/compile"
 	"github.com/matrixorigin/matrixone/pkg/txn/client"
 	"github.com/matrixorigin/matrixone/pkg/txn/trace"
@@ -325,6 +326,42 @@ func TestServiceCloseWithdrawalErrorIsLocallyComplete(t *testing.T) {
 	}
 }
 
+func TestServiceCloseCompletesAfterTransientS3CleanupFailure(t *testing.T) {
+	moruntime.RunTest(t.Name(), func(rt moruntime.Runtime) {
+		withdrawalErr := errors.New("withdrawal failed")
+		hc := &failingWithdrawalHeartbeatClient{testHAKClient: &testHAKClient{}, err: withdrawalErr}
+		mc := clusterservice.NewMOCluster(t.Name(), hc, time.Hour)
+		t.Cleanup(mc.Close)
+		ctrl := gomock.NewController(t)
+		ls := mock_lock.NewMockLockService(ctrl)
+		ls.EXPECT().Close().Return(nil).Times(2)
+		executor := colexec.NewServer(t.Name())
+		var attempts atomic.Int32
+		cleanup := func(context.Context) error {
+			if attempts.Add(1) == 1 {
+				return errors.New("temporary Delete failure")
+			}
+			return nil
+		}
+		sv := &service{
+			cfg: &Config{UUID: t.Name()}, logger: zap.NewNop(), config: util.NewConfigData(nil),
+			stopper: stopper.NewStopper(t.Name()), bootstrapService: &testBootService{},
+			mo: closeErrorMOServer{}, _hakeeperClient: hc, moCluster: mc,
+			server: closeOnlyRPCServer{}, lockService: ls, colexecServer: executor,
+			incrservice: closeOnlyIncrService{onClose: func() {
+				require.NoError(t, executor.RetryUnpublishedS3Cleanup(cleanup))
+			}},
+			viewMetadataAdmissionGeneration: 1,
+		}
+		// Remote withdrawal is diagnostic; it must not hide whether local
+		// teardown advanced past the recovered S3 cleanup and closed its tail.
+		require.ErrorIs(t, sv.Close(), withdrawalErr)
+		require.True(t, sv.CloseComplete())
+		require.Equal(t, int32(2), attempts.Load())
+		require.Equal(t, 1, hc.closed)
+	})
+}
+
 func TestMakeRSSCacheEvictorEvictsMemoryCacheOnly(t *testing.T) {
 	oldMemoryEvictor := evictMemoryCachesToCapacityPercent
 	defer func() {
@@ -1541,6 +1578,55 @@ func TestPipelineAdmissionRejectCancelsRequestOnce(t *testing.T) {
 	require.Equal(t, int32(1), cancelCount.Load())
 }
 
+func TestPipelineBatchAckRunsBeforeIngressReturns(t *testing.T) {
+	started := make(chan struct{})
+	releaseHandler := make(chan struct{})
+	var releaseOnce sync.Once
+	t.Cleanup(func() { releaseOnce.Do(func() { close(releaseHandler) }) })
+	var cancelCount atomic.Int32
+	s := &service{cfg: &Config{}}
+	s.requestHandler = func(
+		_ context.Context,
+		_ string,
+		_ morpc.Message,
+		_ morpc.ClientSession,
+		_ engine.Engine,
+		_ fileservice.FileService,
+		_ lockservice.LockService,
+		_ qclient.QueryClient,
+		_ logservice.CNHAKeeperClient,
+		_ udf.Service,
+		_ client.TxnClient,
+		_ *defines.AutoIncrCacheManager,
+		_ func() morpc.Message,
+	) error {
+		close(started)
+		<-releaseHandler
+		return nil
+	}
+	returned := make(chan error, 1)
+	go func() {
+		returned <- s.handleRequest(context.Background(), morpc.RPCMessage{
+			Message: &pipeline.Message{Cmd: pipeline.Method_PipelineBatchAck, Sid: pipeline.Status_Last},
+			Cancel:  func() { cancelCount.Add(1) },
+		}, 0, nil)
+	}()
+	select {
+	case <-started:
+	case <-time.After(time.Second):
+		t.Fatal("ACK handler did not start")
+	}
+	select {
+	case <-returned:
+		t.Fatal("ingress returned before ACK processing completed")
+	case <-time.After(50 * time.Millisecond):
+	}
+	releaseOnce.Do(func() { close(releaseHandler) })
+	require.NoError(t, <-returned)
+	require.Equal(t, int32(1), cancelCount.Load())
+	require.Zero(t, s.pipelines.counter.Load())
+}
+
 func TestHandleRequestPropagatesConfiguredRPCMaxMessageSize(t *testing.T) {
 	const configuredLimit = 32 * 1024
 	s := &service{cfg: &Config{UUID: t.Name()}}

pkg/objectio/ioutil/sink_pool.go

@@ -194,7 +194,7 @@ func (p *SinkPool) runSyncWorker() {
 		syncStart := time.Now()
 		stats, err := job.fSinker.Sync(r.ctx)
 		atomic.AddInt64(&r.syncNs, int64(time.Since(syncStart)))
-		if err != nil {
+		if err != nil && !objectSyncWasNotStarted(err) {
 			if name := activeFileSinkerObjectName(job.fSinker); name != "" {
 				r.mu.Lock()
 				r.unpublished = append(r.unpublished, name)

pkg/objectio/ioutil/sink_pool_test.go

@@ -16,6 +16,7 @@ package ioutil
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"sync"
 	"sync/atomic"
@@ -151,6 +152,67 @@ func mockFactory(sinkErr, syncErr error) FileSinkerFactory {
 	}
 }
 
+func TestAdmittedFileSinkerReservesBeforeSync(t *testing.T) {
+	base := &mockFileSinker{activeName: "object", syncStart: make(chan struct{})}
+	admitted := &admittedFileSinker{
+		FileSinker: base,
+		reserve: func(name string) error {
+			require.Equal(t, "object", name)
+			return errors.New("cleanup capacity exhausted")
+		},
+	}
+	_, err := admitted.Sync(context.Background())
+	require.ErrorContains(t, err, "cleanup capacity exhausted")
+	select {
+	case <-base.syncStart:
+		t.Fatal("Sync ran despite failed pre-upload admission")
+	default:
+	}
+}
+
+func TestSinkerRejectedAdmissionDoesNotCreateCleanupDebt(t *testing.T) {
+	for _, pipeline := range []bool{false, true} {
+		t.Run(fmt.Sprintf("pipeline=%t", pipeline), func(t *testing.T) {
+			proc := testutil.NewProc(t)
+			fs, err := fileservice.NewMemoryFS("shared", fileservice.DisabledCacheConfig, nil)
+			require.NoError(t, err)
+			attrs, typs, _ := mockSchema(3, -1)
+			base := &mockFileSinker{activeName: "not-written", syncStart: make(chan struct{})}
+			admissionErr := errors.New("cleanup capacity exhausted")
+			opts := []SinkerOption{
+				WithMemorySizeThreshold(1),
+				WithObjectSyncAdmission(func(name string) error {
+					if name != "not-written" {
+						return fmt.Errorf("unexpected object name %q", name)
+					}
+					return admissionErr
+				}, func(string) {}),
+			}
+			if pipeline {
+				opts = append(opts, WithPipelineFlush())
+			}
+			sinker := NewSinker(-1, attrs, typs,
+				func(*mpool.MPool, fileservice.FileService) FileSinker { return base },
+				proc.Mp(), fs, opts...)
+			bat := containers.MockBatch(typs, 8192, -1, nil)
+			err = sinker.Write(context.Background(), containers.ToCNBatch(bat))
+			if err == nil {
+				err = sinker.Sync(context.Background())
+			}
+			require.ErrorIs(t, err, admissionErr)
+			files, cleanupErr := sinker.DeletePersisted(context.Background())
+			require.NoError(t, cleanupErr)
+			require.Empty(t, files, "a rejected upload must not enter cleanup debt")
+			select {
+			case <-base.syncStart:
+				t.Fatal("file Sync started despite rejected admission")
+			default:
+			}
+			_ = sinker.Close()
+		})
+	}
+}
+
 // ---------- SinkPool.Submit tests ----------
 
 func TestSinkPool_SubmitSuccess(t *testing.T) {

pkg/objectio/ioutil/sinker.go

@@ -32,6 +32,7 @@ import (
 	"github.com/matrixorigin/matrixone/pkg/logutil"
 	"github.com/matrixorigin/matrixone/pkg/objectio"
 	"github.com/matrixorigin/matrixone/pkg/objectio/mergeutil"
+	metricv2 "github.com/matrixorigin/matrixone/pkg/util/metric/v2"
 	"github.com/matrixorigin/matrixone/pkg/vm/engine/tae/containers"
 )
 
@@ -107,6 +108,47 @@ func WithChunkedColumnPolicy(policy objectio.ChunkedColumnPolicy) SinkerOption {
 	}
 }
 
+// WithObjectSyncAdmission reserves cleanup capacity before each file Sync.
+// Release runs only after this sinker successfully deletes an object.
+func WithObjectSyncAdmission(reserve func(string) error, release func(string)) SinkerOption {
+	return func(sinker *Sinker) {
+		sinker.config.reserveObject = reserve
+		sinker.config.releaseObject = release
+	}
+}
+
+type admittedFileSinker struct {
+	FileSinker
+	reserve func(string) error
+}
+
+// objectSyncAdmissionError means no object write began. The ordinary Sync
+// error path must retain the name because persistence may be ambiguous.
+type objectSyncAdmissionError struct{ cause error }
+
+func (e *objectSyncAdmissionError) Error() string { return e.cause.Error() }
+func (e *objectSyncAdmissionError) Unwrap() error { return e.cause }
+
+func objectSyncWasNotStarted(err error) bool {
+	var admissionErr *objectSyncAdmissionError
+	return errors.As(err, &admissionErr)
+}
+
+func (s *admittedFileSinker) ActiveObjectName() string {
+	return activeFileSinkerObjectName(s.FileSinker)
+}
+
+func (s *admittedFileSinker) Sync(ctx context.Context) (*objectio.ObjectStats, error) {
+	name := s.ActiveObjectName()
+	if name == "" {
+		return nil, &objectSyncAdmissionError{moerr.NewInvalidStateNoCtx("file sinker has no object name before Sync")}
+	}
+	if err := s.reserve(name); err != nil {
+		return nil, &objectSyncAdmissionError{err}
+	}
+	return s.FileSinker.Sync(ctx)
+}
+
 type FileSinker interface {
 	Sink(context.Context, *batch.Batch) error
 	Sync(context.Context) (*objectio.ObjectStats, error)
@@ -353,6 +395,12 @@ func NewSinker(
 			return fileSinker
 		}
 	}
+	if reserve := sinker.config.reserveObject; reserve != nil {
+		factory := sinker.fSinker.factory
+		sinker.fSinker.factory = func(mp *mpool.MPool, fs fileservice.FileService) FileSinker {
+			return &admittedFileSinker{FileSinker: factory(mp, fs), reserve: reserve}
+		}
+	}
 
 	sinker.fillDefaults()
 	return sinker
@@ -406,6 +454,8 @@ type Sinker struct {
 		tailSizeCap         int
 		offHeap             bool
 		chunkedColumnPolicy objectio.ChunkedColumnPolicy
+		reserveObject       func(string) error
+		releaseObject       func(string)
 	}
 	fSinker struct {
 		executor FileSinker
@@ -493,8 +543,33 @@ func DeleteUnpublishedObjects(
 	fs fileservice.FileService,
 	files ...string,
 ) (int, error) {
+	unique, _, err := deleteUnpublishedObjectBatches(ctx, fs, files)
+	return len(unique), err
+}
+
+// DeleteUnpublishedObjectsWithProgress reports only names in fully confirmed
+// Delete batches. A failed batch can have partial effects, so it and every
+// unattempted batch remain owned. The remaining slice does not retain the
+// backing array of completed names after the caller stores it for retry.
+func DeleteUnpublishedObjectsWithProgress(
+	ctx context.Context,
+	fs fileservice.FileService,
+	files ...string,
+) (completed, remaining []string, err error) {
+	unique, completedCount, err := deleteUnpublishedObjectBatches(ctx, fs, files)
+	if err != nil {
+		return unique[:completedCount], append([]string(nil), unique[completedCount:]...), err
+	}
+	return unique, nil, nil
+}
+
+func deleteUnpublishedObjectBatches(
+	ctx context.Context,
+	fs fileservice.FileService,
+	files []string,
+) (unique []string, completed int, err error) {
 	seen := make(map[string]struct{}, len(files))
-	unique := make([]string, 0, len(files))
+	unique = make([]string, 0, len(files))
 	for _, file := range files {
 		if file == "" {
 			continue
@@ -513,14 +588,15 @@ func DeleteUnpublishedObjects(
 		err := fs.Delete(deleteCtx, unique[start:end]...)
 		cancel()
 		if err != nil && !moerr.IsMoErrCode(err, moerr.ErrFileNotFound) {
-			return len(unique), errors.Join(
+			metricv2.UnpublishedS3DeleteFailuresCounter.Inc()
+			return unique, start, errors.Join(
 				moerr.NewInternalErrorf(
 					ctx, "delete unpublished objects [%d:%d]", start, end),
 				err,
 			)
 		}
 	}
-	return len(unique), nil
+	return unique, len(unique), nil
 }
 
 // DeletePersisted deletes every object that this sinker has persisted, or may
@@ -557,6 +633,11 @@ func (sinker *Sinker) DeletePersisted(ctx context.Context) ([]string, error) {
 	if err != nil {
 		return files, err
 	}
+	if sinker.config.releaseObject != nil {
+		for _, name := range files {
+			sinker.config.releaseObject(name)
+		}
+	}
 
 	sinker.staged.persisted = sinker.staged.persisted[:0]
 	sinker.result.persisted = sinker.result.persisted[:0]
@@ -795,7 +876,7 @@ func (sinker *Sinker) syncFileSinker(ctx context.Context, fSinker FileSinker) er
 	stats, err := fSinker.Sync(ctx)
 	atomic.AddInt64(&sinker.timing.syncNs, int64(time.Since(syncStart)))
 	if err != nil {
-		if name := activeFileSinkerObjectName(fSinker); name != "" {
+		if name := activeFileSinkerObjectName(fSinker); name != "" && !objectSyncWasNotStarted(err) {
 			sinker.staged.unpublished = append(sinker.staged.unpublished, name)
 		}
 		return err