Performance improvements

- Eager async fetch of entries to import
- Drop old tree history (improves situation but it's not a permanent solution, more investigation required) with async iavl tree pruning
- Skip executing import on proposal preparation
- Effectively enforce timeouts through proper context propagation and checking
- Iterate mutable tree with correct iterator when obtaining latest sequence value
- Avoid repeated lock acquisition when importing multiple operations

abciapp/app.go

···
 package abciapp
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"path/filepath"
···
 )
 
 type DIDPLCApplication struct {
+	runnerContext  context.Context
 	plc            plc.PLC
 	tree           *iavl.MutableTree
 	fullyClearTree func() error
···
 func NewDIDPLCApplication(badgerDB *badger.DB, snapshotDirectory string) (*DIDPLCApplication, plc.PLC, func(), error) {
 	treePrefix := []byte{}
 	mkTree := func() *iavl.MutableTree {
-		return iavl.NewMutableTree(badgeradapter.AdaptBadger(badgerDB, treePrefix), 2048, false, iavl.NewNopLogger())
+		return iavl.NewMutableTree(badgeradapter.AdaptBadger(badgerDB, treePrefix), 2048, false, iavl.NewNopLogger(), iavl.AsyncPruningOption(true))
 	}
 
 	tree := mkTree()
···
 	}
 
 	d := &DIDPLCApplication{
+		runnerContext:     context.Background(),
 		tree:              tree,
 		snapshotDirectory: snapshotDirectory,
 		aocsByPLC:         make(map[string]*authoritativeOperationsCache),

abciapp/execution.go

···
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"slices"
 	"time"
 
···
 		if totalSize+len(maybeTx) < int(req.MaxTxBytes)-4096 {
 			// we have space to fit the import transaction
 
-			result, err := processTx(ctx, d.transactionProcessorDependencies(), maybeTx, req.Time, true)
+			// set execute to false to save a lot of time
+			// (we trust that running the import will succeed, so just do bare minimum checks here)
+			result, err := processTx(ctx, d.transactionProcessorDependencies(), maybeTx, req.Time, false)
 			if err != nil {
 				return nil, stacktrace.Propagate(err, "")
 			}
···
 				return &abcitypes.ResponseProcessProposal{Status: abcitypes.ResponseProcessProposal_REJECT}, nil
 			}
 
+			st := time.Now()
 			result, err = finishProcessTx(ctx, d.transactionProcessorDependencies(), processor, tx, req.Time, true)
 			if err != nil {
 				return nil, stacktrace.Propagate(err, "")
 			}
+			fmt.Println("FINISHPROCESSTX TOOK", time.Since(st))
 		}
 
 		// when preparing a proposal, invalid transactions should have been discarded
···
 
 // Commit implements [types.Application].
 func (d *DIDPLCApplication) Commit(context.Context, *abcitypes.RequestCommit) (*abcitypes.ResponseCommit, error) {
-	_, _, err := d.tree.SaveVersion()
+	_, newVersion, err := d.tree.SaveVersion()
 	if err != nil {
 		return nil, stacktrace.Propagate(err, "")
 	}
···
 		}
 	}
 
-	// TODO(later) consider whether we can set some RetainHeight in the response
-	return &abcitypes.ResponseCommit{}, nil
+	minHeightToKeep := max(newVersion-100, 0)
+	minVerToKeep := max(minHeightToKeep-5, 0)
+	if minVerToKeep > 0 {
+		err = d.tree.DeleteVersionsTo(minVerToKeep)
+		if err != nil {
+			return nil, stacktrace.Propagate(err, "")
+		}
+	}
+
+	return &abcitypes.ResponseCommit{
+		// TODO only discard actual blockchain history based on settings
+		//RetainHeight: minHeightToKeep,
+	}, nil
 }
 
 func (d *DIDPLCApplication) transactionProcessorDependencies() TransactionProcessorDependencies {
 	return TransactionProcessorDependencies{
-		plc:       d.plc,
-		tree:      d,
-		aocsByPLC: d.aocsByPLC,
+		runnerContext: d.runnerContext,
+		plc:           d.plc,
+		tree:          d,
+		aocsByPLC:     d.aocsByPLC,
 	}
 }

abciapp/import.go

···
 	"tangled.org/gbl08ma.com/didplcbft/store"
 )
 
+const EagerFetchMaxOps = 10000
+const OpsPerImportTx = 1000
+const OpsPerEagerFetch = 1000
+
 type authoritativeOperationsCache struct {
 	mu sync.Mutex
 
-	plcURL     string
-	operations map[uint64]logEntryWithSeq
+	plcURL               string
+	operations           map[uint64]logEntryWithSeq
+	highestFetchedHeight uint64
 }
 
 type logEntryWithSeq struct {
···
 	Seq uint64 `json:"seq"`
 }
 
-func newAuthoritativeOperationsCache(plc string) *authoritativeOperationsCache {
-	return &authoritativeOperationsCache{
+func newAuthoritativeOperationsCache(ctx context.Context, plc string) *authoritativeOperationsCache {
+	aoc := &authoritativeOperationsCache{
 		plcURL:     plc,
 		operations: make(map[uint64]logEntryWithSeq),
 	}
+
+	go func() {
+		ticker := time.NewTicker(500 * time.Millisecond)
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-ticker.C:
+				aoc.eagerlyFetch(ctx)
+			}
+		}
+	}()
+
+	return aoc
 }
 
-func getOrCreateAuthoritativeOperationsCache(aocsByPLC map[string]*authoritativeOperationsCache, plc string) *authoritativeOperationsCache {
+func getOrCreateAuthoritativeOperationsCache(ctx context.Context, aocsByPLC map[string]*authoritativeOperationsCache, plc string) *authoritativeOperationsCache {
 	aoc, ok := aocsByPLC[plc]
 	if !ok {
-		aoc = newAuthoritativeOperationsCache(plc)
+		aoc = newAuthoritativeOperationsCache(ctx, plc)
 		aocsByPLC[plc] = aoc
 	}
 	return aoc
 }
 
+func (a *authoritativeOperationsCache) eagerlyFetch(ctx context.Context) {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+
+	curOps := len(a.operations)
+	if curOps >= EagerFetchMaxOps {
+		return
+	}
+	_, _ = a.fetchInMutex(ctx, a.highestFetchedHeight, OpsPerEagerFetch)
+}
+
 func (a *authoritativeOperationsCache) dropSeqBelowOrEqual(highestCommittedSeq uint64) {
 	a.mu.Lock()
 	defer a.mu.Unlock()
···
 
 	for _, entry := range entries {
 		a.operations[entry.Seq] = entry
+		a.highestFetchedHeight = max(a.highestFetchedHeight, entry.Seq)
 	}
 	return uint64(len(entries)) < count, nil
 }
···
 		return nil, stacktrace.Propagate(err, "")
 	}
 
-	aoc := getOrCreateAuthoritativeOperationsCache(d.aocsByPLC, plcURL)
+	aoc := getOrCreateAuthoritativeOperationsCache(d.runnerContext, d.aocsByPLC, plcURL)
 
-	entries, err := aoc.get(ctx, cursor, 1000)
+	entries, err := aoc.get(ctx, cursor, OpsPerImportTx)
 	if err != nil {
 		return nil, stacktrace.Propagate(err, "")
 	}

abciapp/tx.go

···
 type TransactionAction string
 
 type TransactionProcessorDependencies struct {
-	plc       plc.PLC
-	tree      plc.TreeProvider // TODO maybe we should move the TreeProvider definition out of the plc package then?
-	aocsByPLC map[string]*authoritativeOperationsCache
+	runnerContext context.Context
+	plc           plc.PLC
+	tree          plc.TreeProvider // TODO maybe we should move the TreeProvider definition out of the plc package then?
+	aocsByPLC     map[string]*authoritativeOperationsCache
 }
 
 type TransactionProcessor func(ctx context.Context, deps TransactionProcessorDependencies, txBytes []byte, atTime time.Time, execute bool) (*processResult, error)

abciapp/tx_import.go

···
 	"net/url"
 	"time"
 
+	"github.com/did-method-plc/go-didplc"
 	cbornode "github.com/ipfs/go-ipld-cbor"
 	"github.com/palantir/stacktrace"
+	"github.com/samber/lo"
 	"tangled.org/gbl08ma.com/didplcbft/plc"
 	"tangled.org/gbl08ma.com/didplcbft/store"
 )
···
 		}, nil
 	}
 
-	aoc := getOrCreateAuthoritativeOperationsCache(deps.aocsByPLC, expectedPlcUrl)
+	aoc := getOrCreateAuthoritativeOperationsCache(deps.runnerContext, deps.aocsByPLC, expectedPlcUrl)
 
 	expectedCursor, err := store.Tree.AuthoritativeImportProgress(roTree)
 	if err != nil {
···
 	}
 
 	if execute {
+		err := deps.plc.ImportOperationsFromAuthoritativeSource(ctx, lo.Map(operations, func(l logEntryWithSeq, _ int) didplc.LogEntry {
+			return l.LogEntry
+		}))
+		if err != nil {
+			return nil, stacktrace.Propagate(err, "")
+		}
+
 		tree, err := deps.tree.MutableTree()
 		if err != nil {
 			return nil, stacktrace.Propagate(err, "")
 		}
 
-		for _, entry := range operations {
-			err := deps.plc.ImportOperationFromAuthoritativeSource(ctx, entry.LogEntry)
-			if err != nil {
-				return nil, stacktrace.Propagate(err, "")
-			}
-		}
 		err = store.Tree.SetAuthoritativeImportProgress(tree, newCursor)
 		if err != nil {
 			return nil, stacktrace.Propagate(err, "")

httpapi/server.go

···
 	"log"
 	"net"
 	"net/http"
+	"net/http/pprof"
 	"slices"
 	"strconv"
 	"strings"
···
 	s.router.HandleFunc("GET /{did}/log/last", s.makeDIDHandler(s.handleGetLastOp))
 	s.router.HandleFunc("GET /{did}/data", s.makeDIDHandler(s.handleGetPLCData))
 	s.router.HandleFunc("GET /export", s.handleExport)
+
+	s.router.HandleFunc("/debug/pprof/", pprof.Index)
+	s.router.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
+	s.router.HandleFunc("/debug/pprof/profile", pprof.Profile)
+	s.router.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
+	s.router.HandleFunc("/debug/pprof/trace", pprof.Trace)
+
+	// Register handlers for specific profiles
+	s.router.Handle("/debug/pprof/heap", pprof.Handler("heap"))
+	s.router.Handle("/debug/pprof/goroutine", pprof.Handler("goroutine"))
+	s.router.Handle("/debug/pprof/block", pprof.Handler("block"))
+	s.router.Handle("/debug/pprof/threadcreate", pprof.Handler("threadcreate"))
 }
 
 // makeDIDHandler creates a wrapper handler that extracts DID from URL path
···
 			return
 		}
 
-		// TODO limit count to 1000 (for debugging it's more useful without limit)
+		if count > 1000 {
+			sendErrorResponse(w, http.StatusBadRequest, "Invalid count parameter")
+			return
+		}
 	}
 
 	afterStr := query.Get("after")

main.go

···
 	badgerDB, err := badger.Open(badger.
 		DefaultOptions(badgerDBPath).
 		WithBlockSize(8 * 1024).
-		WithMemTableSize(256 << 20).
+		WithNumMemtables(3).
+		WithNumLevelZeroTables(3).
 		WithCompression(options.ZSTD))
 	if err != nil {
 		log.Fatalf("Opening badger database: %v", err)
···
 	}()
 
 	if config.PLC.ListenAddress != "" {
-		plcAPIServer, err := httpapi.NewServer(plc, node, config.PLC.ListenAddress, 15*time.Second)
+		plcAPIServer, err := httpapi.NewServer(plc, node, config.PLC.ListenAddress, 30*time.Second)
 		if err != nil {
 			log.Fatalf("Creating PLC API server: %v", err)
 		}

plc/impl.go

···
 	timestamp := syntax.Datetime(at.Format(types.ActualAtprotoDatetimeLayout))
 
 	// TODO set true to false only while importing old ops
-	_, err := plc.validator.Validate(atHeight, timestamp, did, opBytes, true)
+	_, err := plc.validator.Validate(ctx, atHeight, timestamp, did, opBytes, true)
 	if err != nil {
 		return stacktrace.Propagate(err, "operation failed validation")
 	}
···
 	timestamp := syntax.Datetime(t.Format(types.ActualAtprotoDatetimeLayout))
 
 	// TODO set true to false only while importing old ops
-	effects, err := plc.validator.Validate(WorkingTreeVersion, timestamp, did, opBytes, true)
+	effects, err := plc.validator.Validate(ctx, WorkingTreeVersion, timestamp, did, opBytes, true)
 	if err != nil {
 		return stacktrace.Propagate(err, "operation failed validation")
 	}
···
 	return nil
 }
 
+func (plc *plcImpl) ImportOperationsFromAuthoritativeSource(ctx context.Context, newEntries []didplc.LogEntry) error {
+	plc.mu.Lock()
+	defer plc.mu.Unlock()
+
+	tree, err := plc.treeProvider.MutableTree()
+	if err != nil {
+		return stacktrace.Propagate(err, "failed to obtain mutable tree")
+	}
+
+	for _, entry := range newEntries {
+		err := plc.importOp(ctx, tree, entry)
+		if err != nil {
+			return stacktrace.Propagate(err, "")
+		}
+	}
+
+	return nil
+}
+
 func (plc *plcImpl) ImportOperationFromAuthoritativeSource(ctx context.Context, newEntry didplc.LogEntry) error {
 	plc.mu.Lock()
 	defer plc.mu.Unlock()
···
 		return stacktrace.Propagate(err, "failed to obtain mutable tree")
 	}
 
+	return stacktrace.Propagate(plc.importOp(ctx, tree, newEntry), "")
+}
+
+func (plc *plcImpl) importOp(ctx context.Context, tree *iavl.MutableTree, newEntry didplc.LogEntry) error {
 	newCID := newEntry.CID
 	newPrev := newEntry.Operation.AsOperation().PrevCIDStr()
 
 	mostRecentOpIndex := -1
 	indexOfPrev := -1
 	var iteratorErr error
-	for entryIdx, entry := range store.Tree.AuditLogReverseIterator(tree, newEntry.DID, &iteratorErr) {
+	for entryIdx, entry := range store.Tree.AuditLogReverseIterator(ctx, tree, newEntry.DID, &iteratorErr) {
 		entryCID := entry.CID.String()
 		if mostRecentOpIndex == -1 {
 			mostRecentOpIndex = entryIdx
···
 		// there's nothing to do but store the operation, no nullification involved
 		newEntry.Nullified = false
 
-		err = store.Tree.StoreOperation(tree, newEntry, nullifiedEntriesStartingIndex)
+		err := store.Tree.StoreOperation(tree, newEntry, nullifiedEntriesStartingIndex)
 		return stacktrace.Propagate(err, "failed to commit operation")
 	}
 
···
 	}
 
 	newEntry.Nullified = false
-	err = store.Tree.StoreOperation(tree, newEntry, nullifiedEntriesStartingIndex)
+	err := store.Tree.StoreOperation(tree, newEntry, nullifiedEntriesStartingIndex)
 	return stacktrace.Propagate(err, "failed to commit operation")
 }
 
···
 		return didplc.Doc{}, stacktrace.Propagate(err, "failed to obtain immutable tree")
 	}
 
-	l, _, err := store.Tree.AuditLog(tree, did, false)
+	l, _, err := store.Tree.AuditLog(ctx, tree, did, false)
 	if err != nil {
 		return didplc.Doc{}, stacktrace.Propagate(err, "")
 	}
···
 		return nil, stacktrace.Propagate(err, "failed to obtain immutable tree")
 	}
 
-	l, _, err := store.Tree.AuditLog(tree, did, false)
+	l, _, err := store.Tree.AuditLog(ctx, tree, did, false)
 	if err != nil {
 		return nil, stacktrace.Propagate(err, "")
 	}
···
 		return nil, stacktrace.Propagate(err, "failed to obtain immutable tree")
 	}
 
-	l, _, err := store.Tree.AuditLog(tree, did, false)
+	l, _, err := store.Tree.AuditLog(ctx, tree, did, false)
 	if err != nil {
 		return nil, stacktrace.Propagate(err, "")
 	}
···
 		return didplc.OpEnum{}, stacktrace.Propagate(err, "failed to obtain immutable tree")
 	}
 
-	l, _, err := store.Tree.AuditLog(tree, did, false)
+	l, _, err := store.Tree.AuditLog(ctx, tree, did, false)
 	if err != nil {
 		return didplc.OpEnum{}, stacktrace.Propagate(err, "")
 	}
···
 		return didplc.RegularOp{}, stacktrace.Propagate(err, "failed to obtain immutable tree")
 	}
 
-	l, _, err := store.Tree.AuditLog(tree, did, false)
+	l, _, err := store.Tree.AuditLog(ctx, tree, did, false)
 	if err != nil {
 		return didplc.RegularOp{}, stacktrace.Propagate(err, "")
 	}
···
 		return nil, stacktrace.Propagate(err, "failed to obtain immutable tree")
 	}
 
-	entries, err := store.Tree.ExportOperations(tree, after, count)
+	entries, err := store.Tree.ExportOperations(ctx, tree, after, count)
 	return entries, stacktrace.Propagate(err, "")
 }
 
···
 	plc *plcImpl
 }
 
-func (a *inMemoryAuditLogFetcher) AuditLogReverseIterator(atHeight TreeVersion, did string, retErr *error) iter.Seq2[int, types.SequencedLogEntry] {
+func (a *inMemoryAuditLogFetcher) AuditLogReverseIterator(ctx context.Context, atHeight TreeVersion, did string, retErr *error) iter.Seq2[int, types.SequencedLogEntry] {
 	tree, err := a.plc.treeProvider.ImmutableTree(atHeight)
 	if err != nil {
 		*retErr = stacktrace.Propagate(err, "")
 		return func(yield func(int, types.SequencedLogEntry) bool) {}
 	}
 
-	return store.Tree.AuditLogReverseIterator(tree, did, retErr)
+	return store.Tree.AuditLogReverseIterator(ctx, tree, did, retErr)
 }

plc/operation_validator.go

···
 package plc
 
 import (
+	"context"
 	"errors"
 	"iter"
 	"strings"
···
 
 type AuditLogFetcher interface {
 	// AuditLogReverseIterator should return an iterator over the list of log entries for the specified DID, in reverse
-	AuditLogReverseIterator(atHeight TreeVersion, did string, err *error) iter.Seq2[int, types.SequencedLogEntry]
+	AuditLogReverseIterator(ctx context.Context, atHeight TreeVersion, did string, err *error) iter.Seq2[int, types.SequencedLogEntry]
 }
 
 type V0OperationValidator struct {
···
 }
 
 // Validate returns the new complete AuditLog that the DID history would assume if validation passes, and an error if it doesn't pass
-func (v *V0OperationValidator) Validate(atHeight TreeVersion, timestamp syntax.Datetime, expectedDid string, opBytes []byte, laxChecking bool) (OperationEffects, error) {
+func (v *V0OperationValidator) Validate(ctx context.Context, atHeight TreeVersion, timestamp syntax.Datetime, expectedDid string, opBytes []byte, laxChecking bool) (OperationEffects, error) {
 	opEnum, op, err := unmarshalOp(opBytes)
 	if err != nil {
 		return OperationEffects{}, stacktrace.Propagate(errors.Join(ErrMalformedOperation, err), "")
···
 	mostRecentOpIndex := -1
 	indexOfPrev := -1
 	var iteratorErr error
-	for entryIdx, entry := range v.auditLogFetcher.AuditLogReverseIterator(atHeight, expectedDid, &iteratorErr) {
+	for entryIdx, entry := range v.auditLogFetcher.AuditLogReverseIterator(ctx, atHeight, expectedDid, &iteratorErr) {
 		partialLog[entryIdx] = entry
 		if mostRecentOpIndex == -1 {
 			mostRecentOpIndex = entryIdx
···
 	if len(nullifiedEntries) == 0 {
 		// (see prior note on september27Of2023)
 		if !laxChecking && timestamp.Time().After(september29Of2023) {
-			err = v.EnforceOpsRateLimit(atHeight, expectedDid, timestamp.Time())
+			err = v.EnforceOpsRateLimit(ctx, atHeight, expectedDid, timestamp.Time())
 			if err != nil {
 				return OperationEffects{}, stacktrace.Propagate(err, "")
 			}
···
 )
 
 // EnforceOpsRateLimit is ported from the TypeScript enforceOpsRateLimit function, adapted to not require fetching the entire log
-func (v *V0OperationValidator) EnforceOpsRateLimit(atHeight TreeVersion, did string, newOperationTimestamp time.Time) error {
+func (v *V0OperationValidator) EnforceOpsRateLimit(ctx context.Context, atHeight TreeVersion, did string, newOperationTimestamp time.Time) error {
 	hourAgo := newOperationTimestamp.Add(-time.Hour)
 	dayAgo := newOperationTimestamp.Add(-24 * time.Hour)
 	weekAgo := newOperationTimestamp.Add(-7 * 24 * time.Hour)
 
 	var withinHour, withinDay, withinWeek int
 	var err error
-	for _, entry := range v.auditLogFetcher.AuditLogReverseIterator(atHeight, did, &err) {
+	for _, entry := range v.auditLogFetcher.AuditLogReverseIterator(ctx, atHeight, did, &err) {
 		if entry.Nullified {
 			// The typescript implementation operates over a `ops` array which doesn't include nullified ops
 			// (With recovery ops also skipping rate limits, doesn't this leave the PLC vulnerable to the spam of constant recovery operations? TODO investigate)

plc/plc.go

···
 }
 
 type OperationValidator interface {
-	Validate(atHeight TreeVersion, timestamp syntax.Datetime, expectedDid string, opBytes []byte, allowLegacy bool) (OperationEffects, error)
+	Validate(ctx context.Context, atHeight TreeVersion, timestamp syntax.Datetime, expectedDid string, opBytes []byte, allowLegacy bool) (OperationEffects, error)
 }
 
 type PLC interface {
···
 type WritePLC interface {
 	ExecuteOperation(ctx context.Context, timestamp time.Time, did string, opBytes []byte) error
 	ImportOperationFromAuthoritativeSource(ctx context.Context, entry didplc.LogEntry) error
+	ImportOperationsFromAuthoritativeSource(ctx context.Context, entries []didplc.LogEntry) error
 }

store/tree.go

···
 package store
 
 import (
+	"context"
 	"encoding/base32"
 	"encoding/binary"
 	"iter"
···
 var Tree PLCTreeStore = &TreeStore{}
 
 type PLCTreeStore interface {
-	AuditLog(tree ReadOnlyTree, did string, withProof bool) ([]types.SequencedLogEntry, *ics23.CommitmentProof, error)
-	AuditLogReverseIterator(tree ReadOnlyTree, did string, err *error) iter.Seq2[int, types.SequencedLogEntry]
-	ExportOperations(tree ReadOnlyTree, after uint64, count int) ([]types.SequencedLogEntry, error) // passing a count of zero means unlimited
+	AuditLog(ctx context.Context, tree ReadOnlyTree, did string, withProof bool) ([]types.SequencedLogEntry, *ics23.CommitmentProof, error)
+	AuditLogReverseIterator(ctx context.Context, tree ReadOnlyTree, did string, err *error) iter.Seq2[int, types.SequencedLogEntry]
+	ExportOperations(ctx context.Context, tree ReadOnlyTree, after uint64, count int) ([]types.SequencedLogEntry, error) // passing a count of zero means unlimited
 	StoreOperation(tree *iavl.MutableTree, entry didplc.LogEntry, nullifyWithIndexEqualOrGreaterThan mo.Option[int]) error
 	SetOperationCreatedAt(tree *iavl.MutableTree, seqID uint64, createdAt time.Time) error
 
···
 // TreeStore exists just to groups methods nicely
 type TreeStore struct{}
 
-func (t *TreeStore) AuditLog(tree ReadOnlyTree, did string, withProof bool) ([]types.SequencedLogEntry, *ics23.CommitmentProof, error) {
+func (t *TreeStore) AuditLog(ctx context.Context, tree ReadOnlyTree, did string, withProof bool) ([]types.SequencedLogEntry, *ics23.CommitmentProof, error) {
 	proofs := []*ics23.CommitmentProof{}
 
 	didBytes, err := DIDToBytes(did)
···
 
 	logKey := marshalDIDLogKey(didBytes)
 
-	has, err := tree.Has(logKey)
+	logOperations, err := tree.Get(logKey)
 	if err != nil {
 		return nil, nil, stacktrace.Propagate(err, "")
 	}
-
-	var operationKeys [][]byte
-	if has {
-		logOperations, err := tree.Get(logKey)
-		if err != nil {
-			return nil, nil, stacktrace.Propagate(err, "")
-		}
-		operationKeys = make([][]byte, 0, len(logOperations)/8)
-		for seqBytes := range slices.Chunk(logOperations, 8) {
-			operationKeys = append(operationKeys, sequenceBytesToOperationKey(seqBytes))
-		}
+	operationKeys := make([][]byte, 0, len(logOperations)/8)
+	for seqBytes := range slices.Chunk(logOperations, 8) {
+		operationKeys = append(operationKeys, sequenceBytesToOperationKey(seqBytes))
 	}
 
 	if withProof {
···
 
 	logEntries := make([]types.SequencedLogEntry, 0, len(operationKeys))
 	for _, opKey := range operationKeys {
+		select {
+		case <-ctx.Done():
+			return nil, nil, stacktrace.Propagate(ctx.Err(), "")
+		default:
+		}
+
 		operationValue, err := tree.Get(opKey)
 		if err != nil {
 			return nil, nil, stacktrace.Propagate(err, "")
···
 	return logEntries, combinedProof, nil
 }
 
-func (t *TreeStore) AuditLogReverseIterator(tree ReadOnlyTree, did string, retErr *error) iter.Seq2[int, types.SequencedLogEntry] {
+func (t *TreeStore) AuditLogReverseIterator(ctx context.Context, tree ReadOnlyTree, did string, retErr *error) iter.Seq2[int, types.SequencedLogEntry] {
 	return func(yield func(int, types.SequencedLogEntry) bool) {
 		didBytes, err := DIDToBytes(did)
 		if err != nil {
···
 
 		logKey := marshalDIDLogKey(didBytes)
 
-		has, err := tree.Has(logKey)
+		logOperations, err := tree.Get(logKey)
 		if err != nil {
 			*retErr = stacktrace.Propagate(err, "")
 			return
+		}
+		operationKeys := make([][]byte, 0, len(logOperations)/8)
+		for seqBytes := range slices.Chunk(logOperations, 8) {
+			operationKeys = append(operationKeys, sequenceBytesToOperationKey(seqBytes))
 		}
 
-		var operationKeys [][]byte
-		if has {
-			logOperations, err := tree.Get(logKey)
-			if err != nil {
-				*retErr = stacktrace.Propagate(err, "")
+		for i := len(operationKeys) - 1; i >= 0; i-- {
+			select {
+			case <-ctx.Done():
+				*retErr = stacktrace.Propagate(ctx.Err(), "")
 				return
-			}
-			operationKeys = make([][]byte, 0, len(logOperations)/8)
-			for seqBytes := range slices.Chunk(logOperations, 8) {
-				operationKeys = append(operationKeys, sequenceBytesToOperationKey(seqBytes))
+			default:
 			}
-		}
 
-		for i := len(operationKeys) - 1; i >= 0; i-- {
 			opKey := operationKeys[i]
 			operationValue, err := tree.Get(opKey)
 			if err != nil {
···
 	}
 }
 
-func (t *TreeStore) ExportOperations(tree ReadOnlyTree, after uint64, count int) ([]types.SequencedLogEntry, error) {
+func (t *TreeStore) ExportOperations(ctx context.Context, tree ReadOnlyTree, after uint64, count int) ([]types.SequencedLogEntry, error) {
 	// as the name suggests, after is an exclusive lower bound, but our iterators use inclusive lower bounds
 	start := after + 1
 	startKey := marshalOperationKey(start)
···
 	entries := make([]types.SequencedLogEntry, 0, count)
 	var iterErr error
 	tree.IterateRange(startKey, endKey, true, func(operationKey, operationValue []byte) bool {
+		select {
+		case <-ctx.Done():
+			iterErr = stacktrace.Propagate(ctx.Err(), "")
+			return true
+		default:
+		}
+
 		logEntry, err := unmarshalLogEntry(operationKey, operationValue)
 		if err != nil {
 			iterErr = stacktrace.Propagate(err, "")
···
 	opKey := marshalOperationKey(seq)
 	opValue := marshalOperationValue(entry.Nullified, didBytes, opDatetime.Time(), operation)
 
-	_, err = tree.Set(opKey, opValue)
+	updated, err := tree.Set(opKey, opValue)
 	if err != nil {
 		return stacktrace.Propagate(err, "")
 	}
+	if updated {
+		return stacktrace.NewError("expected to be writing to a new operation key but updated instead")
+	}
 
 	logOperations = append(logOperations, opKey[1:9]...)
 	_, err = tree.Set(logKey, logOperations)
···
 
 func getNextSeqID(tree *iavl.MutableTree) (uint64, error) {
 	seq := uint64(0)
-	var err error
-	tree.IterateRange(minOperationKey, maxOperationKey, false, func(key, value []byte) bool {
-		seq, err = unmarshalOperationKey(key)
-		return true
-	})
 
+	itr, err := tree.Iterator(minOperationKey, maxOperationKey, false)
+	if err != nil {
+		return 0, stacktrace.Propagate(err, "")
+	}
+
+	defer itr.Close()
+
+	if itr.Valid() {
+		seq, err = unmarshalOperationKey(itr.Key())
+		if err != nil {
+			return 0, stacktrace.Propagate(err, "")
+		}
+	}
 	return seq + 1, stacktrace.Propagate(err, "")
 }