feat(runtime): Implement branch service (#18475)

Co-authored-by: unknown unknown <unknown@unknown>

runtime/branch.go

@@ -0,0 +1,67 @@
+package runtime
+
+import (
+	"context"
+	"fmt"
+	"os"
+
+	"cosmossdk.io/core/branch"
+	storetypes "cosmossdk.io/store/types"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
+)
+
+var _ branch.Service = BranchService{}
+
+type BranchService struct{}
+
+func (b BranchService) Execute(ctx context.Context, f func(ctx context.Context) error) error {
+	sdkCtx := sdk.UnwrapSDKContext(ctx)
+	branchedCtx, commit := sdkCtx.CacheContext()
+	err := f(branchedCtx)
+	if err != nil {
+		return err
+	}
+	commit()
+	return nil
+}
+
+func (b BranchService) ExecuteWithGasLimit(ctx context.Context, gasLimit uint64, f func(ctx context.Context) error) (gasUsed uint64, err error) {
+	sdkCtx := sdk.UnwrapSDKContext(ctx)
+	branchedCtx, commit := sdkCtx.CacheContext()
+	// create a new gas meter
+	limitedGasMeter := storetypes.NewGasMeter(gasLimit)
+	// apply gas meter with limit to branched context
+	branchedCtx = branchedCtx.WithGasMeter(limitedGasMeter)
+	err = catchOutOfGas(branchedCtx, f)
+	// even before checking the error, we want to get the gas used
+	// and apply it to the original context.
+	gasUsed = limitedGasMeter.GasConsumed()
+	sdkCtx.GasMeter().ConsumeGas(gasUsed, "branch")
+	// in case of errors, do not commit the branched context
+	// return gas used and the error
+	if err != nil {
+		return gasUsed, err
+	}
+	// if no error, commit the branched context
+	// and return gas used and no error
+	commit()
+	return gasUsed, nil
+}
+
+// catchOutOfGas is a helper function to catch out of gas panics and return them as errors.
+func catchOutOfGas(ctx sdk.Context, f func(ctx context.Context) error) (err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			// we immediately check if it's an out of error gas.
+			// if it is not we panic again to propagate it up.
+			if _, ok := r.(storetypes.ErrorOutOfGas); !ok {
+				_, _ = fmt.Fprintf(os.Stderr, "recovered: %#v", r) // log to stderr
+				panic(r)
+			}
+			err = sdkerrors.ErrOutOfGas
+		}
+	}()
+	return f(ctx)
+}

runtime/branch_test.go

@@ -0,0 +1,112 @@
+package runtime
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	storetypes "cosmossdk.io/store/types"
+
+	"github.com/cosmos/cosmos-sdk/testutil"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
+)
+
+func TestBranchService(t *testing.T) {
+	bs := BranchService{}
+	sk := storetypes.NewKVStoreKey("test")
+	tsk := storetypes.NewTransientStoreKey("transient-test")
+	// helper to create a state change
+	doStateChange := func(ctx context.Context) {
+		t.Helper()
+		sdkCtx := sdk.UnwrapSDKContext(ctx)
+		store := sdkCtx.KVStore(sk)
+		store.Set([]byte("key"), []byte("value"))
+	}
+
+	// asserts a state change
+	assertRollback := func(ctx context.Context, shouldRollback bool) {
+		sdkCtx := sdk.UnwrapSDKContext(ctx).WithGasMeter(storetypes.NewInfiniteGasMeter()) // we don't want to consume gas for assertions
+		store := sdkCtx.KVStore(sk)
+		if shouldRollback && store.Has([]byte("key")) {
+			t.Error("expected key to not exist")
+		}
+		if !shouldRollback && !store.Has([]byte("key")) {
+			t.Error("expected key to exist")
+		}
+	}
+
+	t.Run("execute successful", func(t *testing.T) {
+		ctx := testutil.DefaultContext(sk, tsk)
+		err := bs.Execute(ctx, func(ctx context.Context) error {
+			doStateChange(ctx)
+			return nil
+		})
+		require.NoError(t, err)
+		assertRollback(ctx, false)
+	})
+
+	t.Run("execute failed", func(t *testing.T) {
+		ctx := testutil.DefaultContext(sk, tsk)
+		err := bs.Execute(ctx, func(ctx context.Context) error {
+			doStateChange(ctx)
+			return fmt.Errorf("failure")
+		})
+		require.Error(t, err)
+		assertRollback(ctx, true)
+	})
+
+	t.Run("execute with limit successful", func(t *testing.T) {
+		ctx := testutil.DefaultContext(sk, tsk)
+		gasUsed, err := bs.ExecuteWithGasLimit(ctx, 4_000, func(ctx context.Context) error {
+			doStateChange(ctx)
+			return nil
+		})
+		require.NoError(t, err)
+		// assert gas used
+		require.Equal(t, 2240, int(gasUsed))
+		assertRollback(ctx, false)
+		// assert that original context gas was applied
+		require.Equal(t, 2240, int(ctx.GasMeter().GasConsumed()))
+	})
+
+	t.Run("execute with limit failed", func(t *testing.T) {
+		ctx := testutil.DefaultContext(sk, tsk)
+		gasUsed, err := bs.ExecuteWithGasLimit(ctx, 4_000, func(ctx context.Context) error {
+			doStateChange(ctx)
+			return fmt.Errorf("failure")
+		})
+		require.Error(t, err)
+		// assert gas used
+		require.Equal(t, 2240, int(gasUsed))
+		assertRollback(ctx, true)
+		// assert that original context gas was applied
+		require.Equal(t, 2240, int(ctx.GasMeter().GasConsumed()))
+	})
+
+	t.Run("execute with limit out of gas", func(t *testing.T) {
+		ctx := testutil.DefaultContext(sk, tsk)
+		gasUsed, err := bs.ExecuteWithGasLimit(ctx, 2239, func(ctx context.Context) error {
+			doStateChange(ctx)
+			return nil
+		})
+		require.ErrorIs(t, err, sdkerrors.ErrOutOfGas)
+		// assert gas used
+		require.Equal(t, 2240, int(gasUsed))
+		assertRollback(ctx, true)
+		// assert that original context gas was applied
+		require.Equal(t, 2240, int(ctx.GasMeter().GasConsumed()))
+	})
+
+	t.Run("execute with gas limit other panic error", func(t *testing.T) {
+		// ensures other panic errors are not caught by the gas limit panic catcher
+		ctx := testutil.DefaultContext(sk, tsk)
+		require.Panics(t, func() {
+			_, _ = bs.ExecuteWithGasLimit(ctx, 2239, func(ctx context.Context) error {
+				panic("other panic error")
+			})
+		})
+	})
+}