refactor: remove default 1-minute timeout on Database.Close() and set 1-sec in tests (#1458)

Closes #1457

---------

Co-authored-by: Austin Larson <[email protected]>

Author: ARR4N

ffi/firewood.go

@@ -31,7 +31,6 @@ import (
 	"fmt"
 	"runtime"
 	"sync"
-	"time"
 )
 
 const RootLength = C.sizeof_HashKey
@@ -246,18 +245,14 @@ func (db *Database) Revision(root Hash) (*Revision, error) {
 	return rev, nil
 }
 
-// defaultCloseTimeout is the duration by which the [context.Context] passed to
-// [Database.Close] is limited. A minute is arbitrary but well above what is
-// reasonably required, and is chosen simply to avoid permanently blocking.
-var defaultCloseTimeout = time.Minute
-
 // Close releases the memory associated with the Database.
 //
-// This blocks until all outstanding keep-alive handles are disowned. That is,
-// until all Revisions and Proposals created from this Database are either
-// unreachable or one of [Proposal.Commit], [Proposal.Drop], or [Revision.Drop]
-// has been called on them. Unreachable objects will be automatically dropped
-// before Close returns, unless an alternate GC finalizer is set on them.
+// This blocks until all outstanding keep-alive handles are disowned or the
+// [context.Context] is cancelled. That is, until all Revisions and Proposals
+// created from this Database are either unreachable or one of
+// [Proposal.Commit], [Proposal.Drop], or [Revision.Drop] has been called on
+// them. Unreachable objects will be automatically dropped before Close returns,
+// unless an alternate GC finalizer is set on them.
 //
 // This is safe to call if the handle pointer is nil, in which case it does
 // nothing. The pointer will be set to nil after freeing to prevent double free.
@@ -276,12 +271,10 @@ func (db *Database) Close(ctx context.Context) error {
 		close(done)
 	}()
 
-	ctx, cancel := context.WithTimeout(ctx, defaultCloseTimeout)
-	defer cancel()
 	select {
 	case <-done:
 	case <-ctx.Done():
-		return ErrActiveKeepAliveHandles
+		return fmt.Errorf("%w: %w", ctx.Err(), ErrActiveKeepAliveHandles)
 	}
 
 	if err := getErrorFromVoidResult(C.fwd_close_db(db.handle)); err != nil {

ffi/firewood_test.go

@@ -75,6 +75,8 @@ func inferHashingMode(ctx context.Context) (string, error) {
 		return "", err
 	}
 	defer func() {
+		ctx, cancel := context.WithTimeout(ctx, time.Second)
+		defer cancel()
 		_ = db.Close(ctx)
 		_ = os.Remove(dbFile)
 	}()
@@ -137,6 +139,16 @@ func TestMain(m *testing.M) {
 	os.Exit(m.Run())
 }
 
+// oneSecCtx returns `tb.Context()` with a 1-second timeout added. Any existing
+// cancellation on `tb.Context()` is removed, which allows this function to be
+// used inside a `tb.Cleanup()`
+func oneSecCtx(tb testing.TB) context.Context {
+	ctx := context.WithoutCancel(tb.Context())
+	ctx, cancel := context.WithTimeout(ctx, time.Second)
+	tb.Cleanup(cancel)
+	return ctx
+}
+
 func newTestDatabase(t *testing.T, configureFns ...func(*Config)) *Database {
 	t.Helper()
 	r := require.New(t)
@@ -145,7 +157,7 @@ func newTestDatabase(t *testing.T, configureFns ...func(*Config)) *Database {
 	db, err := newDatabase(dbFile, configureFns...)
 	r.NoError(err)
 	t.Cleanup(func() {
-		r.NoError(db.Close(context.Background())) //nolint:usetesting // t.Context() will already be cancelled
+		r.NoError(db.Close(oneSecCtx(t)))
 	})
 	return db
 }
@@ -205,7 +217,7 @@ func TestTruncateDatabase(t *testing.T) {
 	r.NoError(err)
 
 	// Close the database.
-	r.NoError(db.Close(t.Context()))
+	r.NoError(db.Close(oneSecCtx(t)))
 
 	// Reopen the database with truncate enabled.
 	db, err = New(dbFile, config)
@@ -217,7 +229,7 @@ func TestTruncateDatabase(t *testing.T) {
 	expectedHash := stringToHash(t, expectedRoots[emptyKey])
 	r.Equal(expectedHash, hash, "Root hash mismatch after truncation")
 
-	r.NoError(db.Close(t.Context()))
+	r.NoError(db.Close(oneSecCtx(t)))
 }
 
 func TestClosedDatabase(t *testing.T) {
@@ -226,7 +238,7 @@ func TestClosedDatabase(t *testing.T) {
 	db, err := newDatabase(dbFile)
 	r.NoError(err)
 
-	r.NoError(db.Close(t.Context()))
+	r.NoError(db.Close(oneSecCtx(t)))
 
 	_, err = db.Root()
 	r.ErrorIs(err, errDBClosed)
@@ -238,7 +250,7 @@ func TestClosedDatabase(t *testing.T) {
 	r.Empty(root)
 	r.ErrorIs(err, errDBClosed)
 
-	r.NoError(db.Close(t.Context()))
+	r.NoError(db.Close(oneSecCtx(t)))
 }
 
 func keyForTest(i int) []byte {
@@ -1214,7 +1226,7 @@ func TestHandlesFreeImplicitly(t *testing.T) {
 
 	done := make(chan struct{})
 	go func() {
-		require.NoErrorf(t, db.Close(t.Context()), "%T.Close()", db)
+		require.NoErrorf(t, db.Close(oneSecCtx(t)), "%T.Close()", db)
 		close(done)
 	}()
 
@@ -1594,49 +1606,13 @@ func TestCloseWithCancelledContext(t *testing.T) {
 	wg.Wait()
 
 	r.ErrorIs(err, ErrActiveKeepAliveHandles, "Close should return ErrActiveKeepAliveHandles when context is cancelled")
+	r.ErrorIs(err, context.Canceled, "Close error should wrap context.Canceled")
 
 	// Drop the proposal
 	r.NoError(proposal.Drop())
 
 	// Now Close should succeed
-	r.NoError(db.Close(t.Context()))
-}
-
-// TestCloseWithShortTimeout verifies that Database.Close returns
-// ErrActiveKeepAliveHandles when the context times out before handles are dropped.
-func TestCloseWithShortTimeout(t *testing.T) {
-	r := require.New(t)
-	dbFile := filepath.Join(t.TempDir(), "test.db")
-	db, err := newDatabase(dbFile)
-	r.NoError(err)
-
-	// Create a revision to keep a handle active
-	keys, vals := kvForTest(1)
-	root, err := db.Update(keys, vals)
-	r.NoError(err)
-
-	revision, err := db.Revision(root)
-	r.NoError(err)
-
-	// Create a context with a very short timeout (100ms)
-	ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond)
-	defer cancel()
-
-	// Record start time to verify timeout occurred quickly
-	start := time.Now()
-
-	// Close should return ErrActiveKeepAliveHandles after timeout
-	err = db.Close(ctx)
-	elapsed := time.Since(start)
-
-	r.ErrorIs(err, ErrActiveKeepAliveHandles, "Close should return ErrActiveKeepAliveHandles when context times out")
-	r.Less(elapsed, defaultCloseTimeout, "Close should timeout quickly, not wait for default timeout")
-
-	// Drop the revision
-	r.NoError(revision.Drop())
-
-	// Now Close should succeed
-	r.NoError(db.Close(t.Context()))
+	r.NoError(db.Close(oneSecCtx(t)))
 }
 
 // TestCloseWithMultipleActiveHandles verifies that Database.Close returns
@@ -1674,6 +1650,7 @@ func TestCloseWithMultipleActiveHandles(t *testing.T) {
 	// Close should return ErrActiveKeepAliveHandles
 	err = db.Close(ctx)
 	r.ErrorIs(err, ErrActiveKeepAliveHandles, "Close should return ErrActiveKeepAliveHandles with multiple active handles")
+	r.ErrorIs(err, context.Canceled, "Close error should wrap context.Canceled")
 
 	// Drop all handles
 	r.NoError(proposal1.Drop())
@@ -1683,7 +1660,7 @@ func TestCloseWithMultipleActiveHandles(t *testing.T) {
 	r.NoError(revision2.Drop())
 
 	// Now Close should succeed
-	r.NoError(db.Close(t.Context()))
+	r.NoError(db.Close(oneSecCtx(t)))
 }
 
 // TestCloseSucceedsWhenHandlesDroppedInTime verifies that Database.Close succeeds
@@ -1701,28 +1678,17 @@ func TestCloseSucceedsWhenHandlesDroppedInTime(t *testing.T) {
 	proposal2, err := db.Propose(keys[1:2], vals[1:2])
 	r.NoError(err)
 
-	// Create a context with a reasonable timeout
-	ctx, cancel := context.WithTimeout(t.Context(), 500*time.Millisecond)
-	defer cancel()
-
 	// Channel to receive Close result
 	closeDone := make(chan error, 1)
 
 	// Start Close in a goroutine
 	go func() {
-		closeDone <- db.Close(ctx)
+		closeDone <- db.Close(oneSecCtx(t))
 	}()
 
 	// Drop handles after a short delay (before timeout)
 	time.Sleep(100 * time.Millisecond)
 	r.NoError(proposal1.Drop())
 	r.NoError(proposal2.Drop())
-
-	// Close should succeed (not timeout)
-	select {
-	case err := <-closeDone:
-		r.NoError(err, "Close should succeed when handles are dropped before timeout")
-	case <-time.After(defaultCloseTimeout):
-		r.Fail("Close did not complete in time")
-	}
+	r.NoError(<-closeDone, "Close should succeed when handles are dropped before timeout")
 }