feat(ffi): Database.Close() guarantees proposals committed or freed (#1349)

Explicit freeing of proposals propagated through `libevm` (i.e. `geth`)
plumbing has proven difficult when not being committed as they are
simply dropped for the GC to collect. Furthermore, strict ordering of
calls to `Proposal.Drop()` (or `Commit()`) before `Database.Close()` is
required to avoid segfaults. This PR implements a fix for both issues:

1. All new `Proposal`s have a GC finalizer attached, which calls
`Drop()`. This is safe because it is a no-op if called twice or after a
call to `Commit()`.
2. The `Database` has a `sync.WaitGroup` introduced, which tracks all
outstanding proposals. Calls to `Commit()` / `Drop()` decrement the
group counter (only once per `Proposal`).
3. `Database.Close()` waits on the `WaitGroup` before freeing its own
handle, avoiding segfaults.

Assuming that all calls to `Database.Propose()` and `Proposal.Propose()`
occur _before_ the call to `Database.Close()` then this is a correct
usage of `sync.WaitGroup`'s documented requirement for ordering of calls
to `Add()` and `Wait()`.

An integration test demonstrates blocking and eventual return of
`Database.Close()`, specifically due to the unreachability of
un-dropped, un-committed `Proposal`s, resulting in their finalizers
decrementing the `WaitGroup`.

---------

Signed-off-by: Arran Schlosberg <[email protected]>
Co-authored-by: Austin Larson <[email protected]>

Author: 2 people

ffi/firewood.go

@@ -27,9 +27,12 @@ import "C"
 
 import (
 	"bytes"
+	"context"
 	"errors"
 	"fmt"
 	"runtime"
+	"sync"
+	"time"
 )
 
 // These constants are used to identify errors returned by the Firewood Rust FFI.
@@ -49,7 +52,8 @@ type Database struct {
 	// handle is returned and accepted by cgo functions. It MUST be treated as
 	// an opaque value without special meaning.
 	// https://en.wikipedia.org/wiki/Blinkenlights
-	handle *C.DatabaseHandle
+	handle    *C.DatabaseHandle
+	proposals sync.WaitGroup
 }
 
 // Config configures the opening of a [Database].
@@ -139,6 +143,9 @@ func (db *Database) Update(keys, vals [][]byte) ([]byte, error) {
 	return getHashKeyFromHashResult(C.fwd_batch(db.handle, kvp))
 }
 
+// Propose creates a new proposal with the given keys and values. The proposal
+// is not committed until [Proposal.Commit] is called. See [Database.Close] re
+// freeing proposals.
 func (db *Database) Propose(keys, vals [][]byte) (*Proposal, error) {
 	if db.handle == nil {
 		return nil, errDBClosed
@@ -151,8 +158,7 @@ func (db *Database) Propose(keys, vals [][]byte) (*Proposal, error) {
 	if err != nil {
 		return nil, err
 	}
-
-	return getProposalFromProposalResult(C.fwd_propose_on_db(db.handle, kvp), db)
+	return getProposalFromProposalResult(C.fwd_propose_on_db(db.handle, kvp), &db.proposals)
 }
 
 // Get retrieves the value for the given key. It always returns a nil error.
@@ -245,19 +251,43 @@ func (db *Database) Revision(root []byte) (*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 is not safe to call while there are any outstanding Proposals. All proposals
-// must be freed or committed before calling this.
+// 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 is safe to call if the pointer is nil, in which case it does nothing. The
-// pointer will be set to nil after freeing to prevent double free. However, it is
-// not safe to call this method concurrently from multiple goroutines.
-func (db *Database) Close() error {
+// 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.
+// However, it is not safe to call this method concurrently from multiple
+// goroutines.
+func (db *Database) Close(ctx context.Context) error {
 	if db.handle == nil {
 		return nil
 	}
 
+	go runtime.GC()
+
+	done := make(chan struct{})
+	go func() {
+		db.proposals.Wait()
+		close(done)
+	}()
+
+	ctx, cancel := context.WithTimeout(ctx, defaultCloseTimeout)
+	defer cancel()
+	select {
+	case <-done:
+	case <-ctx.Done():
+		return fmt.Errorf("at least one reachable %T neither dropped nor committed", &Proposal{})
+	}
+
 	if err := getErrorFromVoidResult(C.fwd_close_db(db.handle)); err != nil {
 		return fmt.Errorf("unexpected error when closing database: %w", err)
 	}

ffi/firewood_test.go

@@ -5,6 +5,7 @@ package ffi
 
 import (
 	"bytes"
+	"context"
 	"encoding/hex"
 	"errors"
 	"fmt"
@@ -59,14 +60,14 @@ var (
 	expectedRoots map[string]string
 )
 
-func inferHashingMode() (string, error) {
+func inferHashingMode(ctx context.Context) (string, error) {
 	dbFile := filepath.Join(os.TempDir(), "test.db")
-	db, closeDB, err := newDatabase(dbFile)
+	db, err := newDatabase(dbFile)
 	if err != nil {
 		return "", err
 	}
 	defer func() {
-		_ = closeDB()
+		_ = db.Close(ctx)
 		_ = os.Remove(dbFile)
 	}()
 
@@ -111,7 +112,7 @@ func TestMain(m *testing.M) {
 	// Otherwise, infer the hash mode from an empty database.
 	hashMode := os.Getenv("TEST_FIREWOOD_HASH_MODE")
 	if hashMode == "" {
-		inferredHashMode, err := inferHashingMode()
+		inferredHashMode, err := inferHashingMode(context.Background())
 		if err != nil {
 			fmt.Fprintf(os.Stderr, "failed to infer hash mode %v\n", err)
 			os.Exit(1)
@@ -133,15 +134,15 @@ func newTestDatabase(t *testing.T, configureFns ...func(*Config)) *Database {
 	r := require.New(t)
 
 	dbFile := filepath.Join(t.TempDir(), "test.db")
-	db, closeDB, err := newDatabase(dbFile, configureFns...)
+	db, err := newDatabase(dbFile, configureFns...)
 	r.NoError(err)
 	t.Cleanup(func() {
-		r.NoError(closeDB())
+		r.NoError(db.Close(context.Background())) //nolint:usetesting // t.Context() will already be cancelled
 	})
 	return db
 }
 
-func newDatabase(dbFile string, configureFns ...func(*Config)) (*Database, func() error, error) {
+func newDatabase(dbFile string, configureFns ...func(*Config)) (*Database, error) {
 	conf := DefaultConfig()
 	conf.Truncate = true // in tests, we use filepath.Join, which creates an empty file
 	for _, fn := range configureFns {
@@ -150,9 +151,9 @@ func newDatabase(dbFile string, configureFns ...func(*Config)) (*Database, func(
 
 	f, err := New(dbFile, conf)
 	if err != nil {
-		return nil, nil, fmt.Errorf("failed to create new database at filepath %q: %w", dbFile, err)
+		return nil, fmt.Errorf("failed to create new database at filepath %q: %w", dbFile, err)
 	}
-	return f, f.Close, nil
+	return f, nil
 }
 
 func TestUpdateSingleKV(t *testing.T) {
@@ -196,7 +197,7 @@ func TestTruncateDatabase(t *testing.T) {
 	r.NoError(err)
 
 	// Close the database.
-	r.NoError(db.Close())
+	r.NoError(db.Close(t.Context()))
 
 	// Reopen the database with truncate enabled.
 	db, err = New(dbFile, config)
@@ -210,16 +211,16 @@ func TestTruncateDatabase(t *testing.T) {
 	r.NoError(err)
 	r.Equal(expectedHash, hash, "Root hash mismatch after truncation")
 
-	r.NoError(db.Close())
+	r.NoError(db.Close(t.Context()))
 }
 
 func TestClosedDatabase(t *testing.T) {
 	r := require.New(t)
 	dbFile := filepath.Join(t.TempDir(), "test.db")
-	db, _, err := newDatabase(dbFile)
+	db, err := newDatabase(dbFile)
 	r.NoError(err)
 
-	r.NoError(db.Close())
+	r.NoError(db.Close(t.Context()))
 
 	_, err = db.Root()
 	r.ErrorIs(err, errDBClosed)
@@ -231,7 +232,7 @@ func TestClosedDatabase(t *testing.T) {
 	r.Empty(root)
 	r.ErrorIs(err, errDBClosed)
 
-	r.NoError(db.Close())
+	r.NoError(db.Close(t.Context()))
 }
 
 func keyForTest(i int) []byte {
@@ -1171,6 +1172,63 @@ func TestGetFromRootParallel(t *testing.T) {
 	}
 }
 
+func TestProposalHandlesFreed(t *testing.T) {
+	t.Parallel()
+
+	db, err := newDatabase(filepath.Join(t.TempDir(), "test_GC_drops_proposal.db"))
+	require.NoError(t, err)
+
+	// These MUST NOT be committed nor dropped as they demonstrate that the GC
+	// finalizer does it for us.
+	p0, err := db.Propose(kvForTest(1))
+	require.NoErrorf(t, err, "%T.Propose(...)", db)
+	p1, err := p0.Propose(kvForTest(1))
+	require.NoErrorf(t, err, "%T.Propose(...)", p0)
+
+	// Demonstrates that explicit [Proposal.Commit] and [Proposal.Drop] calls
+	// are sufficient to unblock [Database.Close].
+	var keep []*Proposal
+	for name, free := range map[string](func(*Proposal) error){
+		"Commit": (*Proposal).Commit,
+		"Drop":   (*Proposal).Drop,
+	} {
+		p, err := db.Propose(kvForTest(1))
+		require.NoErrorf(t, err, "%T.Propose(...)", db)
+		require.NoErrorf(t, free(p), "%T.%s()", p, name)
+		keep = append(keep, p)
+	}
+
+	done := make(chan struct{})
+	go func() {
+		require.NoErrorf(t, db.Close(t.Context()), "%T.Close()", db)
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		t.Errorf("%T.Close() returned with undropped %T", db, p0) //nolint:forbidigo // Use of require is impossible without a hack like require.False(true)
+	case <-time.After(300 * time.Millisecond):
+		// TODO(arr4n) use `synctest` package when at Go 1.25
+	}
+
+	runtime.KeepAlive(p0)
+	runtime.KeepAlive(p1)
+	p0 = nil
+	p1 = nil //nolint:ineffassign // Makes the value unreachable, allowing the finalizer to call Drop()
+
+	// In practice there's no need to call [runtime.GC] if [Database.Close] is
+	// called after all proposals are unreachable, as it does it itself.
+	runtime.GC()
+	// Note that [Database.Close] waits for outstanding proposals, so this would
+	// block permanently if the unreachability of `p0` and `p1` didn't result in
+	// their [Proposal.Drop] methods being called.
+	<-done
+
+	for _, p := range keep {
+		runtime.KeepAlive(p)
+	}
+}
+
 type kvIter interface {
 	SetBatchSize(int)
 	Next() bool

ffi/kvbackend.go

@@ -38,9 +38,6 @@ type kVBackend interface {
 	// commits happen on a rolling basis.
 	// Length of the root slice is guaranteed to be common.HashLength.
 	Commit(root []byte) error
-
-	// Close closes the backend and releases all held resources.
-	Close() error
 }
 
 // Prefetch is a no-op since we don't need to prefetch for Firewood.