feat(ffi): add keepalive struct to own waitgroup logic (#1437)

Expanding on Arran and Austin's work, I added the
`databaseKeepAliveHandle` to manage the WaitGroup logic for keeping the
database alive while various objects needed.

In the future, the `RangeProof` object will also use this keep-alive
handle, but only during certain stages of its lifecycle. Because of
this, the disown method on the keep-alive handle is capable of being
called on the nil receiver without causing a panic.

Author: demosdemon

ffi/firewood.go

@@ -39,8 +39,9 @@ const RootLength = C.sizeof_HashKey
 type Hash [RootLength]byte
 
 var (
-	EmptyRoot   Hash
-	errDBClosed = errors.New("firewood database already closed")
+	EmptyRoot                 Hash
+	errDBClosed               = errors.New("firewood database already closed")
+	ErrActiveKeepAliveHandles = errors.New("cannot close database with active keep-alive handles")
 )
 
 // A Database is a handle to a Firewood database.
@@ -252,10 +253,11 @@ var defaultCloseTimeout = time.Minute
 
 // Close releases the memory associated with the Database.
 //
-// This blocks until all outstanding Proposals are either unreachable or one of
-// [Proposal.Commit] or [Proposal.Drop] has been called on them. Unreachable
-// proposals 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. 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.
@@ -279,7 +281,7 @@ func (db *Database) Close(ctx context.Context) error {
 	select {
 	case <-done:
 	case <-ctx.Done():
-		return fmt.Errorf("at least one reachable %T neither dropped nor committed", &Proposal{})
+		return ErrActiveKeepAliveHandles
 	}
 
 	if err := getErrorFromVoidResult(C.fwd_close_db(db.handle)); err != nil {

ffi/firewood_test.go

@@ -1566,3 +1566,163 @@ func TestNilVsEmptyValue(t *testing.T) {
 	r.NotNil(got, "key4 should exist")
 	r.Empty(got, "key4 should have empty value")
 }
+
+// TestCloseWithCancelledContext verifies that Database.Close returns
+// ErrActiveKeepAliveHandles when the context is cancelled before handles are dropped.
+func TestCloseWithCancelledContext(t *testing.T) {
+	r := require.New(t)
+	dbFile := filepath.Join(t.TempDir(), "test.db")
+	db, err := newDatabase(dbFile)
+	r.NoError(err)
+
+	// Create a proposal to keep a handle active
+	keys, vals := kvForTest(1)
+	proposal, err := db.Propose(keys, vals)
+	r.NoError(err)
+
+	ctx, cancel := context.WithCancel(t.Context())
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+
+		err = db.Close(ctx)
+	}()
+
+	cancel()
+	wg.Wait()
+
+	r.ErrorIs(err, ErrActiveKeepAliveHandles, "Close should return ErrActiveKeepAliveHandles when context is cancelled")
+
+	// 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()))
+}
+
+// TestCloseWithMultipleActiveHandles verifies that Database.Close returns
+// ErrActiveKeepAliveHandles when multiple handles are active and context is cancelled.
+func TestCloseWithMultipleActiveHandles(t *testing.T) {
+	r := require.New(t)
+	dbFile := filepath.Join(t.TempDir(), "test.db")
+	db, err := newDatabase(dbFile)
+	r.NoError(err)
+
+	// Create multiple proposals
+	keys1, vals1 := kvForTest(3)
+	proposal1, err := db.Propose(keys1[:1], vals1[:1])
+	r.NoError(err)
+	proposal2, err := db.Propose(keys1[1:2], vals1[1:2])
+	r.NoError(err)
+	proposal3, err := db.Propose(keys1[2:3], vals1[2:3])
+	r.NoError(err)
+
+	// Create multiple revisions
+	root1, err := db.Update([][]byte{keyForTest(10)}, [][]byte{valForTest(10)})
+	r.NoError(err)
+	root2, err := db.Update([][]byte{keyForTest(20)}, [][]byte{valForTest(20)})
+	r.NoError(err)
+
+	revision1, err := db.Revision(root1)
+	r.NoError(err)
+	revision2, err := db.Revision(root2)
+	r.NoError(err)
+
+	// Create a cancelled context
+	ctx, cancel := context.WithCancel(t.Context())
+	cancel()
+
+	// Close should return ErrActiveKeepAliveHandles
+	err = db.Close(ctx)
+	r.ErrorIs(err, ErrActiveKeepAliveHandles, "Close should return ErrActiveKeepAliveHandles with multiple active handles")
+
+	// Drop all handles
+	r.NoError(proposal1.Drop())
+	r.NoError(proposal2.Drop())
+	r.NoError(proposal3.Drop())
+	r.NoError(revision1.Drop())
+	r.NoError(revision2.Drop())
+
+	// Now Close should succeed
+	r.NoError(db.Close(t.Context()))
+}
+
+// TestCloseSucceedsWhenHandlesDroppedInTime verifies that Database.Close succeeds
+// when all handles are dropped before the context timeout.
+func TestCloseSucceedsWhenHandlesDroppedInTime(t *testing.T) {
+	r := require.New(t)
+	dbFile := filepath.Join(t.TempDir(), "test.db")
+	db, err := newDatabase(dbFile)
+	r.NoError(err)
+
+	// Create two active proposals
+	keys, vals := kvForTest(2)
+	proposal1, err := db.Propose(keys[:1], vals[:1])
+	r.NoError(err)
+	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)
+	}()
+
+	// 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")
+	}
+}

ffi/keepalive.go

@@ -0,0 +1,61 @@
+// Copyright (C) 2025, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE.md for licensing terms.
+
+package ffi
+
+import "sync"
+
+// databaseKeepAliveHandle is added to types that hold a lease on the database
+// to ensure it is not closed while those types are still in use.
+//
+// This is necessary to prevent use-after-free bugs where a type holding a
+// reference to the database outlives the database itself. Even attempting to
+// free those objects after the database has been closed will lead to undefined
+// behavior, as a part of the underling Rust object will have already been freed.
+type databaseKeepAliveHandle struct {
+	mu sync.Mutex
+	// [Database.Close] blocks on this WaitGroup, which is set and incremented
+	// by [newKeepAliveHandle], and decremented by
+	// [databaseKeepAliveHandle.disown].
+	outstandingHandles *sync.WaitGroup
+}
+
+// init initializes the keep-alive handle to track a new outstanding handle.
+func (h *databaseKeepAliveHandle) init(wg *sync.WaitGroup) {
+	// lock not necessary today, but will be necessary in the future for types
+	// that initialize the handle at some point after construction (#1429).
+	h.mu.Lock()
+	defer h.mu.Unlock()
+
+	if h.outstandingHandles != nil {
+		// setting the finalizer twice will also panic, so we're panicking
+		// early to provide better context
+		panic("keep-alive handle already initialized")
+	}
+
+	h.outstandingHandles = wg
+	h.outstandingHandles.Add(1)
+}
+
+// disown indicates that the object owning this handle is no longer keeping the
+// database alive. If [attemptDisown] returns an error, disowning will only occur
+// if [disownEvenOnErr] is true.
+//
+// This method is safe to call multiple times; subsequent calls after the first
+// will continue to invoke [attemptDisown] but will not decrement the wait group
+// unless [databaseKeepAliveHandle.init] was called again in the meantime.
+func (h *databaseKeepAliveHandle) disown(disownEvenOnErr bool, attemptDisown func() error) error {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+
+	err := attemptDisown()
+
+	if (err == nil || disownEvenOnErr) && h.outstandingHandles != nil {
+		h.outstandingHandles.Done()
+		// prevent calling `Done` multiple times if disown is called again, which
+		// may happen when the finalizer runs after an explicit call to Drop or Commit.
+		h.outstandingHandles = nil
+	}
+
+	return err
+}

ffi/proposal.go

@@ -20,6 +20,15 @@ import (
 
 var errDroppedProposal = errors.New("proposal already dropped")
 
+// Proposal represents a set of proposed changes to be committed to the database.
+// Proposals are created via [Database.Propose] or [Proposal.Propose], and must be
+// either committed with [Proposal.Commit] or released with [Proposal.Drop].
+//
+// Proposals must be committed or dropped before the associated database is
+// closed. A finalizer is set on each Proposal to ensure that Drop is called
+// when the Proposal is garbage collected, but relying on finalizers is not
+// recommended. Failing to commit or drop a proposal before the database is
+// closed will cause it to block or fail.
 type Proposal struct {
 	// handle is an opaque pointer to the proposal within Firewood. It should be
 	// passed to the C FFI functions that operate on proposals
@@ -29,14 +38,14 @@ type Proposal struct {
 	// Calls to `C.fwd_commit_proposal` and `C.fwd_free_proposal` will invalidate
 	// this handle, so it should not be used after those calls.
 	handle *C.ProposalHandle
-	disown sync.Mutex
-	// [Database.Close] blocks on this WaitGroup, which is incremented by
-	// [getProposalFromProposalResult], and decremented by either
-	// [Proposal.Commit] or [Proposal.Drop] (when the handle is disowned).
-	outstandingHandles *sync.WaitGroup
 
-	// The proposal root hash.
+	// root is the root hash of the proposal and the expected root hash after commit.
 	root Hash
+
+	// keepAliveHandle is used to keep the database alive while this proposal is
+	// in use. It is initialized when the proposal is created and disowned after
+	// [Proposal.Commit] or [Proposal.Drop] is called.
+	keepAliveHandle databaseKeepAliveHandle
 }
 
 // Root retrieves the root hash of the proposal.
@@ -93,37 +102,25 @@ func (p *Proposal) Propose(keys, vals [][]byte) (*Proposal, error) {
 	if err != nil {
 		return nil, err
 	}
-	return getProposalFromProposalResult(C.fwd_propose_on_proposal(p.handle, kvp), p.outstandingHandles)
-}
-
-// disownHandle is the common path of [Proposal.Commit] and [Proposal.Drop], the
-// `fn` argument defining the method-specific behaviour.
-func (p *Proposal) disownHandle(fn func(*C.ProposalHandle) error, disownEvenOnErr bool) error {
-	p.disown.Lock()
-	defer p.disown.Unlock()
-
-	if p.handle == nil {
-		return errDroppedProposal
-	}
-	err := fn(p.handle)
-	if disownEvenOnErr || err == nil {
-		p.handle = nil
-		p.outstandingHandles.Done()
-	}
-	return err
+	return getProposalFromProposalResult(C.fwd_propose_on_proposal(p.handle, kvp), p.keepAliveHandle.outstandingHandles)
 }
 
 // Commit commits the proposal and returns any errors.
 //
 // The proposal handle is no longer valid after this call, but the root
-// hash can still be retrieved using Root().
+// hash can still be retrieved using [Proposal.Root].
 func (p *Proposal) Commit() error {
-	return p.disownHandle(commitProposal, true)
-}
+	return p.keepAliveHandle.disown(true /* evenOnError */, func() error {
+		if p.handle == nil {
+			return errDroppedProposal
+		}
+
+		_, err := getHashKeyFromHashResult(C.fwd_commit_proposal(p.handle))
 
-func commitProposal(h *C.ProposalHandle) error {
-	_, err := getHashKeyFromHashResult(C.fwd_commit_proposal(h))
-	return err
+		p.handle = nil
+
+		return err
+	})
 }
 
 // Drop releases the memory associated with the Proposal.
@@ -132,33 +129,34 @@ func commitProposal(h *C.ProposalHandle) error {
 //
 // The pointer will be set to nil after freeing to prevent double free.
 func (p *Proposal) Drop() error {
-	if err := p.disownHandle(dropProposal, false); err != nil && err != errDroppedProposal {
-		return err
-	}
-	return nil
-}
+	return p.keepAliveHandle.disown(false /* evenOnError */, func() error {
+		if p.handle == nil {
+			return nil
+		}
 
-func dropProposal(h *C.ProposalHandle) error {
-	if err := getErrorFromVoidResult(C.fwd_free_proposal(h)); err != nil {
-		return fmt.Errorf("%w: %w", errFreeingValue, err)
-	}
-	return nil
+		if err := getErrorFromVoidResult(C.fwd_free_proposal(p.handle)); err != nil {
+			return fmt.Errorf("%w: %w", errFreeingValue, err)
+		}
+
+		p.handle = nil
+
+		return nil
+	})
 }
 
 // getProposalFromProposalResult converts a C.ProposalResult to a Proposal or error.
-func getProposalFromProposalResult(result C.ProposalResult, outstandingHandles *sync.WaitGroup) (*Proposal, error) {
+func getProposalFromProposalResult(result C.ProposalResult, wg *sync.WaitGroup) (*Proposal, error) {
 	switch result.tag {
 	case C.ProposalResult_NullHandlePointer:
 		return nil, errDBClosed
 	case C.ProposalResult_Ok:
 		body := (*C.ProposalResult_Ok_Body)(unsafe.Pointer(&result.anon0))
 		hashKey := *(*Hash)(unsafe.Pointer(&body.root_hash._0))
 		proposal := &Proposal{
-			handle:             body.handle,
-			root:               hashKey,
-			outstandingHandles: outstandingHandles,
+			handle: body.handle,
+			root:   hashKey,
 		}
-		outstandingHandles.Add(1)
+		proposal.keepAliveHandle.init(wg)
 		runtime.SetFinalizer(proposal, (*Proposal).Drop)
 		return proposal, nil
 	case C.ProposalResult_Err: