fix: keep admitted mutation RPCs alive until their result arrives

Signed-off-by: Feng Ruohang <rh@vonng.com>

cmd/storage-rest-client.go

@@ -412,7 +412,7 @@ func (client *storageRESTClient) WriteMetadata(ctx context.Context, origvolume,
 		return err
 	}
 	// As with RenameData, admitted metadata writes must finish before unlock.
-	ctx = context.WithoutCancel(ctx)
+	ctx = grid.WithoutDefaultTimeout(context.WithoutCancel(ctx))
 
 	_, err := storageWriteMetadataRPC.Call(ctx, client.gridConn, &MetadataHandlerParams{
 		DiskID:     *client.diskID.Load(),
@@ -428,7 +428,7 @@ func (client *storageRESTClient) UpdateMetadata(ctx context.Context, volume, pat
 	if err := ctx.Err(); err != nil {
 		return err
 	}
-	ctx = context.WithoutCancel(ctx)
+	ctx = grid.WithoutDefaultTimeout(context.WithoutCancel(ctx))
 
 	_, err := storageUpdateMetadataRPC.Call(ctx, client.gridConn, &MetadataHandlerParams{
 		DiskID:     *client.diskID.Load(),
@@ -443,7 +443,7 @@ func (client *storageRESTClient) UpdateMetadata(ctx context.Context, volume, pat
 func (client *storageRESTClient) DeleteVersion(ctx context.Context, volume, path string, fi FileInfo, forceDelMarker bool, opts DeleteOptions) (err error) {
 	if opts.UndoWrite {
 		// Rollback is part of the commit and must finish before unlocking.
-		ctx = context.WithoutCancel(ctx)
+		ctx = grid.WithoutDefaultTimeout(context.WithoutCancel(ctx))
 	} else {
 		var cancel context.CancelFunc
 		ctx, cancel = context.WithTimeout(ctx, globalDriveConfig.GetMaxTimeout())
@@ -504,7 +504,7 @@ func (client *storageRESTClient) RenameData(ctx context.Context, srcVolume, srcP
 	// The remote handler cannot cancel an admitted RenameData. Keep waiting
 	// for its result so renameData can roll back every completed mutation
 	// before the caller cleans staging files and releases the namespace lock.
-	ctx = context.WithoutCancel(ctx)
+	ctx = grid.WithoutDefaultTimeout(context.WithoutCancel(ctx))
 	params := RenameDataHandlerParams{
 		DiskID:    *client.diskID.Load(),
 		SrcVolume: srcVolume,

internal/grid/connection.go

@@ -344,7 +344,8 @@ func (c *Subroute) newMuxClient(ctx context.Context) (*muxClient, error) {
 
 // Request allows to do a single remote request.
 // 'req' will not be used after the call and caller can reuse.
-// If no deadline is set on ctx, a 1-minute deadline will be added.
+// If no deadline is set on ctx, a 1-minute deadline will be added unless
+// WithoutDefaultTimeout was requested.
 func (c *Connection) Request(ctx context.Context, h HandlerID, req []byte) ([]byte, error) {
 	if !h.valid() {
 		return nil, ErrUnknownHandler
@@ -370,7 +371,8 @@ func (c *Connection) Request(ctx context.Context, h HandlerID, req []byte) ([]by
 
 // Request allows to do a single remote request.
 // 'req' will not be used after the call and caller can reuse.
-// If no deadline is set on ctx, a 1-minute deadline will be added.
+// If no deadline is set on ctx, a 1-minute deadline will be added unless
+// WithoutDefaultTimeout was requested.
 func (c *Subroute) Request(ctx context.Context, h HandlerID, req []byte) ([]byte, error) {
 	if !h.valid() {
 		return nil, ErrUnknownHandler

internal/grid/handlers.go

@@ -550,7 +550,8 @@ type Requester interface {
 
 // Call the remote with the request and return the response.
 // The response should be returned with PutResponse when no error.
-// If no deadline is set, a 1-minute deadline is added.
+// If no deadline is set, a 1-minute deadline is added unless
+// WithoutDefaultTimeout was requested.
 func (h *SingleHandler[Req, Resp]) Call(ctx context.Context, c Requester, req Req) (resp Resp, err error) {
 	if c == nil {
 		if h.ignoreNilConn {

internal/grid/muxclient.go

@@ -59,6 +59,16 @@ type Response struct {
 	Err error
 }
 
+type withoutDefaultTimeoutKey struct{}
+
+// WithoutDefaultTimeout disables the implicit single-request timeout.
+// Explicit context cancellation and connection loss still end the request.
+// Use this for admitted mutations that must return their actual result before
+// the caller releases its lock.
+func WithoutDefaultTimeout(ctx context.Context) context.Context {
+	return context.WithValue(ctx, withoutDefaultTimeoutKey{}, true)
+}
+
 func newMuxClient(ctx context.Context, muxID uint64, parent *Connection) *muxClient {
 	ctx, cancelFn := context.WithCancelCause(ctx)
 	return &muxClient{
@@ -101,8 +111,10 @@ func (m *muxClient) roundtrip(h HandlerID, req []byte) ([]byte, error) {
 	m.respMu.Unlock()
 	ctx := m.ctx
 
-	// Add deadline if none.
-	if msg.DeadlineMS == 0 {
+	// Admitted storage mutations must not time out while the handler can
+	// still commit. A zero wire deadline also keeps their late response.
+	noDefaultTimeout, _ := ctx.Value(withoutDefaultTimeoutKey{}).(bool)
+	if msg.DeadlineMS == 0 && !noDefaultTimeout {
 		msg.DeadlineMS = uint32(defaultSingleRequestTimeout / time.Millisecond)
 		var cancel context.CancelFunc
 		ctx, cancel = context.WithTimeout(ctx, defaultSingleRequestTimeout)

internal/grid/muxclient_timeout_test.go

@@ -0,0 +1,109 @@
+// Copyright (c) 2026 Feng Ruohang
+//
+// This file is part of Silo Object Storage.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see <http://www.gnu.org/licenses/>.
+
+package grid
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"testing/synctest"
+	"time"
+)
+
+func TestSingleMutationWaitsBeyondDefaultTimeout(t *testing.T) {
+	for _, admitted := range []bool{false, true} {
+		t.Run(map[bool]string{false: "ordinary", true: "admitted"}[admitted], func(t *testing.T) {
+			synctest.Test(t, func(t *testing.T) {
+				ctx := t.Context()
+				if admitted {
+					ctx = WithoutDefaultTimeout(ctx)
+				}
+				conn := &Connection{outQueue: make(chan []byte, 1)}
+				client := newMuxClient(ctx, 1, conn)
+				defer client.cancelFn(context.Canceled)
+				done := make(chan Response, 1)
+				go func() {
+					msg, err := client.roundtrip(handlerTest, nil)
+					done <- Response{Msg: msg, Err: err}
+				}()
+				wire := <-conn.outQueue
+				var request message
+				if _, err := request.UnmarshalMsg(wire); err != nil {
+					t.Fatal(err)
+				}
+				PutByteBuffer(wire)
+				if admitted && request.DeadlineMS != 0 {
+					t.Errorf("admitted mutation tells the server to discard a late result after %dms", request.DeadlineMS)
+				}
+				time.Sleep(defaultSingleRequestTimeout + time.Second)
+				synctest.Wait()
+				if !admitted {
+					select {
+					case got := <-done:
+						if !errors.Is(got.Err, context.DeadlineExceeded) {
+							t.Fatalf("ordinary timeout: %v", got.Err)
+						}
+					default:
+						t.Fatal("ordinary request lost its default timeout")
+					}
+					return
+				}
+				select {
+				case got := <-done:
+					t.Fatalf("mutation returned before its commit completed: %v", got.Err)
+				default:
+				}
+				client.response(0, Response{Msg: []byte("committed")})
+				got := <-done
+				if got.Err != nil || string(got.Msg) != "committed" {
+					t.Fatalf("lost late commit result: %q, %v", got.Msg, got.Err)
+				}
+			})
+		})
+	}
+}
+
+func TestWithoutDefaultTimeoutStillEndsOnDisconnectOrCancellation(t *testing.T) {
+	for _, disconnect := range []bool{false, true} {
+		t.Run(map[bool]string{false: "explicit-cancel", true: "disconnect"}[disconnect], func(t *testing.T) {
+			synctest.Test(t, func(t *testing.T) {
+				ctx, cancel := context.WithCancel(t.Context())
+				defer cancel()
+				conn := &Connection{outQueue: make(chan []byte, 1)}
+				client := newMuxClient(WithoutDefaultTimeout(ctx), 1, conn)
+				defer client.cancelFn(context.Canceled)
+				done := make(chan error, 1)
+				go func() {
+					_, err := client.roundtrip(handlerTest, nil)
+					done <- err
+				}()
+				PutByteBuffer(<-conn.outQueue)
+				want := context.Canceled
+				if disconnect {
+					client.close()
+					want = ErrDisconnected
+				} else {
+					cancel()
+				}
+				if got := <-done; !errors.Is(got, want) {
+					t.Fatalf("request did not terminate correctly: want %v, got %v", want, got)
+				}
+			})
+		})
+	}
+}