Merge pull request ethereum#25 from flashbots/feature/onfly-sessions

Add on-the-fly sessions for specific endpoints
allnil · Mar 15, 2024 · 3d1bbf1 · 3d1bbf1
2 parents 287a107 + ad08846
commit 3d1bbf1
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 24 deletions.
diff --git a/suave/builder/session_manager.go b/suave/builder/session_manager.go
@@ -70,20 +70,7 @@ func (s *SessionManager) TxPool() *txpool.TxPool {
 	return s.pool
 }
 
-// NewSession creates a new builder session and returns the session id
-func (s *SessionManager) NewSession(ctx context.Context, args *api.BuildBlockArgs) (string, error) {
-	if args == nil {
-		return "", fmt.Errorf("args cannot be nil")
-	}
-	// Wait for session to become available
-	select {
-	case <-s.sem:
-		s.sessionsLock.Lock()
-		defer s.sessionsLock.Unlock()
-	case <-ctx.Done():
-		return "", ctx.Err()
-	}
-
+func (s *SessionManager) newBuilder(args *api.BuildBlockArgs) (*miner.Builder, error) {
 	builderCfg := &miner.BuilderConfig{
 		ChainConfig: s.blockchain.Config(),
 		Engine:      s.blockchain.Engine(),
@@ -100,11 +87,33 @@ func (s *SessionManager) NewSession(ctx context.Context, args *api.BuildBlockArg
 		Slot:           args.Slot,
 	}
 
-	id := uuid.New().String()[:7]
 	session, err := miner.NewBuilder(builderCfg, builderArgs)
+	if err != nil {
+		return nil, err
+	}
+	return session, nil
+}
+
+// NewSession creates a new builder session and returns the session id
+func (s *SessionManager) NewSession(ctx context.Context, args *api.BuildBlockArgs) (string, error) {
+	if args == nil {
+		return "", fmt.Errorf("args cannot be nil")
+	}
+	// Wait for session to become available
+	select {
+	case <-s.sem:
+		s.sessionsLock.Lock()
+		defer s.sessionsLock.Unlock()
+	case <-ctx.Done():
+		return "", ctx.Err()
+	}
+
+	session, err := s.newBuilder(args)
 	if err != nil {
 		return "", err
 	}
+
+	id := uuid.New().String()[:7]
 	s.sessions[id] = session
 
 	// start session timer
@@ -127,7 +136,11 @@ func (s *SessionManager) NewSession(ctx context.Context, args *api.BuildBlockArg
 	return id, nil
 }
 
-func (s *SessionManager) getSession(sessionId string) (*miner.Builder, error) {
+func (s *SessionManager) getSession(sessionId string, allowOnTheFlySession bool) (*miner.Builder, error) {
+	if sessionId == "" {
+		return s.newBuilder(&api.BuildBlockArgs{})
+	}
+
 	s.sessionsLock.RLock()
 	defer s.sessionsLock.RUnlock()
 
@@ -143,31 +156,31 @@ func (s *SessionManager) getSession(sessionId string) (*miner.Builder, error) {
 }
 
 func (s *SessionManager) AddTransaction(sessionId string, tx *types.Transaction) (*api.SimulateTransactionResult, error) {
-	builder, err := s.getSession(sessionId)
+	builder, err := s.getSession(sessionId, true)
 	if err != nil {
 		return nil, err
 	}
 	return builder.AddTransaction(tx)
 }
 
 func (s *SessionManager) AddTransactions(sessionId string, txs types.Transactions) ([]*api.SimulateTransactionResult, error) {
-	builder, err := s.getSession(sessionId)
+	builder, err := s.getSession(sessionId, true)
 	if err != nil {
 		return nil, err
 	}
 	return builder.AddTransactions(txs)
 }
 
 func (s *SessionManager) AddBundles(sessionId string, bundles []*api.Bundle) ([]*api.SimulateBundleResult, error) {
-	builder, err := s.getSession(sessionId)
+	builder, err := s.getSession(sessionId, true)
 	if err != nil {
 		return nil, err
 	}
 	return builder.AddBundles(bundles)
 }
 
 func (s *SessionManager) BuildBlock(sessionId string) error {
-	builder, err := s.getSession(sessionId)
+	builder, err := s.getSession(sessionId, false)
 	if err != nil {
 		return err
 	}
@@ -176,7 +189,7 @@ func (s *SessionManager) BuildBlock(sessionId string) error {
 }
 
 func (s *SessionManager) Bid(sessionId string, blsPubKey phase0.BLSPubKey) (*api.SubmitBlockRequest, error) {
-	builder, err := s.getSession(sessionId)
+	builder, err := s.getSession(sessionId, false)
 	if err != nil {
 		return nil, err
 	}

diff --git a/suave/builder/session_manager_test.go b/suave/builder/session_manager_test.go
@@ -33,7 +33,7 @@ func TestSessionManager_SessionTimeout(t *testing.T) {
 
 	time.Sleep(1 * time.Second)
 
-	_, err = mngr.getSession(id)
+	_, err = mngr.getSession(id, false)
 	require.Error(t, err)
 }
 
@@ -89,7 +89,7 @@ func TestSessionManager_SessionRefresh(t *testing.T) {
 	for i := 0; i < 5; i++ {
 		time.Sleep(250 * time.Millisecond)
 
-		_, err = mngr.getSession(id)
+		_, err = mngr.getSession(id, false)
 		require.NoError(t, err)
 	}
 
@@ -98,7 +98,7 @@ func TestSessionManager_SessionRefresh(t *testing.T) {
 
 	time.Sleep(1 * time.Second)
 
-	_, err = mngr.getSession(id)
+	_, err = mngr.getSession(id, false)
 	require.Error(t, err)
 }
 
@@ -114,6 +114,11 @@ func TestSessionManager_StartSession(t *testing.T) {
 	receipt, err := mngr.AddTransaction(id, txn)
 	require.NoError(t, err)
 	require.NotNil(t, receipt)
+
+	// test that you can simulate the transaction on the fly
+	receipt2, err := mngr.AddTransaction("", txn)
+	require.NoError(t, err)
+	require.Equal(t, receipt, receipt2)
 }
 
 func newSessionManager(t *testing.T, cfg *Config) (*SessionManager, *testBackend) {