From 01356c72bb9ffa1f08dd16dc9b0e654ca0033c4b Mon Sep 17 00:00:00 2001
From: Georgi <georgi.goranov@zama.ai>
Date: Wed, 18 Dec 2024 18:19:31 +0200
Subject: [PATCH] feat(late-commit): make late-commit configurable, disabled by
 default

---
 fhevm-engine/fhevm-go-native/fhevm/api.go | 52 +++++++++++++++++++----
 1 file changed, 44 insertions(+), 8 deletions(-)

diff --git a/fhevm-engine/fhevm-go-native/fhevm/api.go b/fhevm-engine/fhevm-go-native/fhevm/api.go
index ecac202..e673453 100644
--- a/fhevm-engine/fhevm-go-native/fhevm/api.go
+++ b/fhevm-engine/fhevm-go-native/fhevm/api.go
@@ -8,6 +8,7 @@ import (
 	"math/big"
 	"os"
 	"sort"
+	"strconv"
 	"strings"
 	"sync"
 	"time"
@@ -200,6 +201,18 @@ type ApiImpl struct {
 	executorUrl            string
 	contractStorageAddress common.Address
 	cache                  *CiphertextCache
+
+	// The offset from the current block number for committing the FHE computations.
+	// If set to 0, the computations are committed in the current block.
+	commitBlockOffset uint8
+}
+
+// Notify the executor that there is work available
+func (s *ApiImpl) notifyWorkAvailable() {
+	select {
+	case s.cache.workAvailableChan <- true:
+	default:
+	}
 }
 
 type SessionImpl struct {
@@ -260,6 +273,7 @@ type SessionComputationStore struct {
 	blockNumber            int64
 	cache                  *CiphertextCache
 	contractStorageAddress common.Address
+	commitBlockOffset      uint8
 }
 
 type EvmStorageComputationStore struct {
@@ -279,6 +293,7 @@ func (executorApi *ApiImpl) CreateSession(blockNumber int64) ExecutorSession {
 			blockNumber:            blockNumber,
 			cache:                  executorApi.cache,
 			contractStorageAddress: executorApi.contractStorageAddress,
+			commitBlockOffset:      executorApi.commitBlockOffset,
 		},
 	}
 }
@@ -405,6 +420,20 @@ func (sessionApi *SessionImpl) Commit(blockNumber int64, storage ChainStorageApi
 		return err
 	}
 
+	// Compute pending computations
+	if sessionApi.apiImpl.commitBlockOffset == 0 {
+		// Late commit is disabled, send compute gRPC request and waits for it to finish
+		err := executorProcessPendingComputations(sessionApi.apiImpl)
+		if err != nil {
+			log.Error("Executor failed", "block", blockNumber, "error", err)
+			return err
+		}
+	} else {
+		// Signal the executor thread that work is ready.
+		sessionApi.apiImpl.notifyWorkAvailable()
+	}
+
+	// Flush the computation results to the state
 	err = sessionApi.apiImpl.flushFheResultsToState(blockNumber, storage)
 	if err != nil {
 		return err
@@ -502,7 +531,7 @@ func (dbApi *SessionComputationStore) InsertComputation(computation ComputationT
 		// hardcode late commit for now to be 5 blocks from current block
 		// in future we can implement dynamic compute, if user pays more
 		// he can have faster commit
-		computation.CommitBlockId = dbApi.blockNumber + 5
+		computation.CommitBlockId = dbApi.blockNumber + int64(dbApi.commitBlockOffset)
 		dbApi.inserts = append(dbApi.inserts, computation)
 		log.Info("Insert computation",
 			"inserts count", len(dbApi.inserts), "computation", computation)
@@ -752,12 +781,6 @@ func (dbApi *EvmStorageComputationStore) InsertComputationBatch(evmStorage Chain
 
 	}
 
-	// notify about work available
-	select {
-	case dbApi.cache.workAvailableChan <- true:
-	default:
-	}
-
 	return nil
 }
 
@@ -996,6 +1019,16 @@ func InitExecutor() (ExecutorApi, error) {
 	// pick hardcoded value in the beginning, we can change later
 	storageAddress := common.HexToAddress("0x0000000000000000000000000000000000000070")
 
+	commitBlockOffset := uint8(0)
+	offset, hasOffset := os.LookupEnv("FHEVM_COMMIT_BLOCK_OFFSET")
+	if hasOffset {
+		parsedOffset, err := strconv.ParseUint(offset, 10, 8)
+		if err != nil {
+			logger("init").Crit("Invalid FHEVM_COMMIT_BLOCK_OFFSET", "error", err.Error())
+		}
+		commitBlockOffset = uint8(parsedOffset)
+	}
+
 	logger("init").Info("FHEVM initialized",
 		"Executor addr", executorUrl,
 		"FHEVM contract", contractAddr,
@@ -1018,10 +1051,13 @@ func InitExecutor() (ExecutorApi, error) {
 		contractStorageAddress: storageAddress,
 		executorUrl:            executorUrl,
 		cache:                  cache,
+		commitBlockOffset:      uint8(commitBlockOffset),
 	}
 
 	// run executor worker in the background
-	go executorWorkerThread(apiImpl)
+	if commitBlockOffset > 0 {
+		go executorWorkerThread(apiImpl)
+	}
 
 	return apiImpl, nil
 }