diff --git a/ffi/firewood.go b/ffi/firewood.go index 06f3188e28..6e2233ec79 100644 --- a/ffi/firewood.go +++ b/ffi/firewood.go @@ -214,6 +214,17 @@ func (db *Database) Root() ([]byte, error) { return bytes, err } +func (db *Database) LatestRevision() (*Revision, error) { + root, err := db.Root() + if err != nil { + return nil, err + } + if bytes.Equal(root, EmptyRoot) { + return nil, errRevisionNotFound + } + return db.Revision(root) +} + // Revision returns a historical revision of the database. func (db *Database) Revision(root []byte) (*Revision, error) { if root == nil || len(root) != RootLength { diff --git a/ffi/firewood_test.go b/ffi/firewood_test.go index 76356952a9..88523b10fe 100644 --- a/ffi/firewood_test.go +++ b/ffi/firewood_test.go @@ -1136,3 +1136,144 @@ func TestGetFromRootParallel(t *testing.T) { r.NoError(err, "Parallel operation failed") } } + +func assertIteratorYields(r *require.Assertions, it *Iterator, keys [][]byte, vals [][]byte) { + i := 0 + for ; it.Next(); i += 1 { + r.Equal(keys[i], it.Key()) + r.Equal(vals[i], it.Value()) + } + r.NoError(it.Err()) + r.Equal(len(keys), i) +} + +// Tests that basic iterator functionality works +func TestIter(t *testing.T) { + r := require.New(t) + db := newTestDatabase(t) + + keys, vals := kvForTest(10) + _, err := db.Update(keys, vals) + r.NoError(err) + + rev, err := db.LatestRevision() + r.NoError(err) + it, err := rev.Iter(nil) + r.NoError(err) + t.Cleanup(func() { + r.NoError(it.Drop()) + r.NoError(rev.Drop()) + }) + + assertIteratorYields(r, it, keys, vals) +} + +// Tests that iterators on different roots work fine +func TestIterOnRoot(t *testing.T) { + r := require.New(t) + db := newTestDatabase(t) + + // Commit 10 key-value pairs. + keys, vals := kvForTest(20) + keys = keys[:10] + vals1, vals2 := vals[:10], vals[10:] + + firstRoot, err := db.Update(keys, vals1) + r.NoError(err) + + // we use the same keys, but update the values + secondRoot, err := db.Update(keys, vals2) + r.NoError(err) + + r1, err := db.Revision(firstRoot) + r.NoError(err) + h1, err := r1.Iter(nil) + r.NoError(err) + t.Cleanup(func() { + r.NoError(h1.Drop()) + r.NoError(r1.Drop()) + }) + + r2, err := db.Revision(secondRoot) + r.NoError(err) + h2, err := r2.Iter(nil) + r.NoError(err) + t.Cleanup(func() { + r.NoError(h2.Drop()) + r.NoError(r2.Drop()) + }) + + assertIteratorYields(r, h1, keys, vals1) + assertIteratorYields(r, h2, keys, vals2) +} + +// Tests that basic iterator functionality works for proposal +func TestIterOnProposal(t *testing.T) { + r := require.New(t) + db := newTestDatabase(t) + + keys, vals := kvForTest(10) + p, err := db.Propose(keys, vals) + r.NoError(err) + + it, err := p.Iter(nil) + r.NoError(err) + t.Cleanup(func() { + r.NoError(it.Drop()) + }) + + assertIteratorYields(r, it, keys, vals) +} + +// Tests that the iterator still works after proposal is committed +func TestIterAfterProposalCommit(t *testing.T) { + r := require.New(t) + db := newTestDatabase(t) + + keys, vals := kvForTest(10) + p, err := db.Propose(keys, vals) + r.NoError(err) + + it, err := p.Iter(nil) + r.NoError(err) + t.Cleanup(func() { + r.NoError(it.Drop()) + }) + + err = p.Commit() + r.NoError(err) + + // iterate after commit + // because iterator hangs on the nodestore reference of proposal + // the nodestore won't be dropped until we drop the iterator + assertIteratorYields(r, it, keys, vals) +} + +// Tests that the iterator on latest revision works properly after a proposal commit +func TestIterUpdate(t *testing.T) { + r := require.New(t) + db := newTestDatabase(t) + + keys, vals := kvForTest(10) + _, err := db.Update(keys, vals) + r.NoError(err) + + // get an iterator on latest revision + rev, err := db.LatestRevision() + r.NoError(err) + it, err := rev.Iter(nil) + r.NoError(err) + t.Cleanup(func() { + r.NoError(it.Drop()) + r.NoError(rev.Drop()) + }) + + // update the database + keys2, vals2 := kvForTest(10) + _, err = db.Update(keys2, vals2) + r.NoError(err) + + // iterate after commit + // because iterator is fixed on the revision hash, it should return the initial values + assertIteratorYields(r, it, keys, vals) +} diff --git a/ffi/iterator.go b/ffi/iterator.go new file mode 100644 index 0000000000..9bf5fa190e --- /dev/null +++ b/ffi/iterator.go @@ -0,0 +1,92 @@ +// Copyright (C) 2025, Ava Labs, Inc. All rights reserved. +// See the file LICENSE.md for licensing terms. + +package ffi + +// #include +// #include "firewood.h" +import "C" + +import ( + "fmt" + "unsafe" +) + +type Iterator struct { + // handle is an opaque pointer to the iterator within Firewood. It should be + // passed to the C FFI functions that operate on iterators + // + // It is not safe to call these methods with a nil handle. + handle *C.IteratorHandle + + // currentKey is the current key retrieved from the iterator + currentKey []byte + // currentVal is the current value retrieved from the iterator + currentVal []byte + // err is the error from the iterator, if any + err error +} + +// Next proceeds to the next item on the iterator, and returns true +// if succeeded and there is a pair available. +// The new pair could be retrieved with Key and Value methods. +func (it *Iterator) Next() bool { + kv, e := getKeyValueFromKeyValueResult(C.fwd_iter_next(it.handle)) + it.err = e + if kv == nil || e != nil { + return false + } + k, v, e := kv.Consume() + it.currentKey = k + it.currentVal = v + it.err = e + return e == nil +} + +// Key returns the key of the current pair +func (it *Iterator) Key() []byte { + if (it.currentKey == nil && it.currentVal == nil) || it.err != nil { + return nil + } + return it.currentKey +} + +// Value returns the value of the current pair +func (it *Iterator) Value() []byte { + if (it.currentKey == nil && it.currentVal == nil) || it.err != nil { + return nil + } + return it.currentVal +} + +// Err returns the error if Next failed +func (it *Iterator) Err() error { + return it.err +} + +// Drop drops the iterator and releases the resources +func (it *Iterator) Drop() error { + if it.handle != nil { + return getErrorFromVoidResult(C.fwd_free_iterator(it.handle)) + } + return nil +} + +// getIteratorFromIteratorResult converts a C.IteratorResult to an Iterator or error. +func getIteratorFromIteratorResult(result C.IteratorResult) (*Iterator, error) { + switch result.tag { + case C.IteratorResult_NullHandlePointer: + return nil, errDBClosed + case C.IteratorResult_Ok: + body := (*C.IteratorResult_Ok_Body)(unsafe.Pointer(&result.anon0)) + proposal := &Iterator{ + handle: body.handle, + } + return proposal, nil + case C.IteratorResult_Err: + err := newOwnedBytes(*(*C.OwnedBytes)(unsafe.Pointer(&result.anon0))).intoError() + return nil, err + default: + return nil, fmt.Errorf("unknown C.IteratorResult tag: %d", result.tag) + } +} diff --git a/ffi/memory.go b/ffi/memory.go index 441964ca9f..0e2c632895 100644 --- a/ffi/memory.go +++ b/ffi/memory.go @@ -309,6 +309,56 @@ func getValueFromValueResult(result C.ValueResult) ([]byte, error) { } } +type ownedKeyValue struct { + key *ownedBytes + value *ownedBytes +} + +func (kv *ownedKeyValue) Consume() ([]byte, []byte, error) { + key := kv.key.CopiedBytes() + if err := kv.key.Free(); err != nil { + return nil, nil, fmt.Errorf("%w: %w", errFreeingValue, err) + } + value := kv.value.CopiedBytes() + if err := kv.value.Free(); err != nil { + return nil, nil, fmt.Errorf("%w: %w", errFreeingValue, err) + } + return key, value, nil +} + +// newOwnedKeyValue creates a ownedKeyValue from a C.OwnedKeyValuePair. +// +// The caller is responsible for calling Free() on the returned ownedKeyValue +// when it is no longer needed otherwise memory will leak. +func newOwnedKeyValue(owned C.OwnedKeyValuePair) *ownedKeyValue { + return &ownedKeyValue{ + key: newOwnedBytes(owned.key), + value: newOwnedBytes(owned.value), + } +} + +// getKeyValueFromKeyValueResult converts a C.KeyValueResult to a key value pair or error. +// +// It returns nil, nil if the result is None. +// It returns a *ownedKeyValue, nil if the result is Some. +// It returns an error if the result is an error. +func getKeyValueFromKeyValueResult(result C.KeyValueResult) (*ownedKeyValue, error) { + switch result.tag { + case C.KeyValueResult_NullHandlePointer: + return nil, errDBClosed + case C.KeyValueResult_None: + return nil, nil + case C.KeyValueResult_Some: + ownedKvp := newOwnedKeyValue(*(*C.OwnedKeyValuePair)(unsafe.Pointer(&result.anon0))) + return ownedKvp, nil + case C.ValueResult_Err: + err := newOwnedBytes(*(*C.OwnedBytes)(unsafe.Pointer(&result.anon0))).intoError() + return nil, err + default: + return nil, fmt.Errorf("unknown C.KeyValueResult tag: %d", result.tag) + } +} + // getDatabaseFromHandleResult converts a C.HandleResult to a Database or error. // // If the C.HandleResult is an error, it returns an error instead of a Database. diff --git a/ffi/proposal.go b/ffi/proposal.go index 285c14f5ec..56ed0ad1a4 100644 --- a/ffi/proposal.go +++ b/ffi/proposal.go @@ -58,6 +58,21 @@ func (p *Proposal) Get(key []byte) ([]byte, error) { return getValueFromValueResult(C.fwd_get_from_proposal(p.handle, newBorrowedBytes(key, &pinner))) } +// Iter creates and iterator starting from the provided key on proposal. +// pass empty slice to start from beginning +func (p *Proposal) Iter(key []byte) (*Iterator, error) { + if p.handle == nil { + return nil, errDBClosed + } + + var pinner runtime.Pinner + defer pinner.Unpin() + + itResult := C.fwd_iter_on_proposal(p.handle, newBorrowedBytes(key, &pinner)) + + return getIteratorFromIteratorResult(itResult) +} + // Propose creates a new proposal with the given keys and values. // The proposal is not committed until Commit is called. func (p *Proposal) Propose(keys, vals [][]byte) (*Proposal, error) { diff --git a/ffi/revision.go b/ffi/revision.go index e884a43a88..0fc600aa42 100644 --- a/ffi/revision.go +++ b/ffi/revision.go @@ -53,6 +53,21 @@ func (r *Revision) Get(key []byte) ([]byte, error) { )) } +// Iter creates an iterator starting from the provided key on revision. +// pass empty slice to start from beginning +func (r *Revision) Iter(key []byte) (*Iterator, error) { + if r.handle == nil { + return nil, errDroppedRevision + } + + var pinner runtime.Pinner + defer pinner.Unpin() + + itResult := C.fwd_iter_on_revision(r.handle, newBorrowedBytes(key, &pinner)) + + return getIteratorFromIteratorResult(itResult) +} + // Drop releases the resources backed by the revision handle. // // It is safe to call Drop multiple times; subsequent calls after the first are no-ops. diff --git a/ffi/src/iterator.rs b/ffi/src/iterator.rs index 487caf1cb6..94e6a5d644 100644 --- a/ffi/src/iterator.rs +++ b/ffi/src/iterator.rs @@ -1,35 +1,35 @@ // Copyright (C) 2025, Ava Labs, Inc. All rights reserved. // See the file LICENSE.md for licensing terms. -use std::ops::{Deref, DerefMut}; - use derive_where::derive_where; +use firewood::merkle; +use firewood::v2::api; use firewood::v2::api::BoxKeyValueIter; /// An opaque wrapper around a [`BoxKeyValueIter`]. +#[derive(Default)] #[derive_where(Debug)] #[derive_where(skip_inner)] -pub struct IteratorHandle<'view>(BoxKeyValueIter<'view>); +pub struct IteratorHandle<'view>(Option>); impl<'view> From> for IteratorHandle<'view> { fn from(value: BoxKeyValueIter<'view>) -> Self { - IteratorHandle(value) + IteratorHandle(Some(value)) } } -impl<'view> Deref for IteratorHandle<'view> { - type Target = BoxKeyValueIter<'view>; - - fn deref(&self) -> &Self::Target { - &self.0 - } -} +impl Iterator for IteratorHandle<'_> { + type Item = Result<(merkle::Key, merkle::Value), api::Error>; -impl DerefMut for IteratorHandle<'_> { - fn deref_mut(&mut self) -> &mut Self::Target { - &mut self.0 + fn next(&mut self) -> Option { + let out = self.0.as_mut()?.next(); + if out.is_none() { + // iterator exhausted; drop it so the NodeStore can be released + self.0 = None; + } + out } } -#[derive(Debug)] +#[derive(Debug, Default)] pub struct CreateIteratorResult<'db>(pub IteratorHandle<'db>); diff --git a/ffi/src/lib.rs b/ffi/src/lib.rs index 4d54bdb7e9..77515bae39 100644 --- a/ffi/src/lib.rs +++ b/ffi/src/lib.rs @@ -192,7 +192,7 @@ pub unsafe extern "C" fn fwd_iter_on_proposal<'p>( /// #[unsafe(no_mangle)] pub unsafe extern "C" fn fwd_iter_next(handle: Option<&mut IteratorHandle<'_>>) -> KeyValueResult { - invoke_with_handle(handle, |it| it.next()) + invoke_with_handle(handle, Iterator::next) } /// Consumes the [`IteratorHandle`], destroys the iterator, and frees the memory. diff --git a/ffi/src/value/results.rs b/ffi/src/value/results.rs index a351f79d87..ee94c97e37 100644 --- a/ffi/src/value/results.rs +++ b/ffi/src/value/results.rs @@ -355,8 +355,8 @@ pub enum KeyValueResult { Err(OwnedBytes), } -impl From>> for KeyValueResult { - fn from(value: Option>) -> Self { +impl From>> for KeyValueResult { + fn from(value: Option>) -> Self { match value { Some(value) => match value { Ok(value) => KeyValueResult::Some(value.into()),