From 2dfac08e80743022601d3d8928202dbceb54b491 Mon Sep 17 00:00:00 2001 From: Louis Thibault Date: Mon, 22 Jan 2024 16:22:37 -0500 Subject: [PATCH] Hack: use a semaphore channel to limit the number of open sessions. This is a hack because we're relying on session timeout behavior. Recommend ref-counting API that releases refs when block-height changes. --- suave/builder/api/api_server.go | 4 +-- suave/builder/api/api_test.go | 6 ++-- suave/builder/session_manager.go | 37 +++++++++++++++++++---- suave/builder/session_manager_test.go | 42 +++++++++++++++++++++++++-- 4 files changed, 75 insertions(+), 14 deletions(-) diff --git a/suave/builder/api/api_server.go b/suave/builder/api/api_server.go index 9ed45465fe..4b8336ed06 100644 --- a/suave/builder/api/api_server.go +++ b/suave/builder/api/api_server.go @@ -8,7 +8,7 @@ import ( // SessionManager is the backend that manages the session state of the builder API. type SessionManager interface { - NewSession() (string, error) + NewSession(context.Context) (string, error) AddTransaction(sessionId string, tx *types.Transaction) (*types.SimulateTransactionResult, error) } @@ -24,7 +24,7 @@ type Server struct { } func (s *Server) NewSession(ctx context.Context) (string, error) { - return s.sessionMngr.NewSession() + return s.sessionMngr.NewSession(ctx) } func (s *Server) AddTransaction(ctx context.Context, sessionId string, tx *types.Transaction) (*types.SimulateTransactionResult, error) { diff --git a/suave/builder/api/api_test.go b/suave/builder/api/api_test.go index 51122eef02..f9d795033c 100644 --- a/suave/builder/api/api_test.go +++ b/suave/builder/api/api_test.go @@ -30,10 +30,10 @@ func TestAPI(t *testing.T) { type nullSessionManager struct{} -func (n *nullSessionManager) NewSession() (string, error) { - return "1", nil +func (nullSessionManager) NewSession(ctx context.Context) (string, error) { + return "1", ctx.Err() } -func (n *nullSessionManager) AddTransaction(sessionId string, tx *types.Transaction) (*types.SimulateTransactionResult, error) { +func (nullSessionManager) AddTransaction(sessionId string, tx *types.Transaction) (*types.SimulateTransactionResult, error) { return &types.SimulateTransactionResult{Logs: []*types.SimulatedLog{}}, nil } diff --git a/suave/builder/session_manager.go b/suave/builder/session_manager.go index 2ad4d7376b..905241a5bd 100644 --- a/suave/builder/session_manager.go +++ b/suave/builder/session_manager.go @@ -1,6 +1,7 @@ package builder import ( + "context" "fmt" "math/big" "sync" @@ -31,11 +32,13 @@ type blockchain interface { } type Config struct { - GasCeil uint64 - SessionIdleTimeout time.Duration + GasCeil uint64 + SessionIdleTimeout time.Duration + MaxConcurrentSessions int } type SessionManager struct { + sem chan struct{} sessions map[string]*builder sessionTimers map[string]*time.Timer sessionsLock sync.RWMutex @@ -50,8 +53,17 @@ func NewSessionManager(blockchain blockchain, config *Config) *SessionManager { if config.SessionIdleTimeout == 0 { config.SessionIdleTimeout = 5 * time.Second } + if config.MaxConcurrentSessions <= 0 { + config.MaxConcurrentSessions = 16 // chosen arbitrarily + } + + sem := make(chan struct{}, config.MaxConcurrentSessions) + for len(sem) < cap(sem) { + sem <- struct{}{} // fill 'er up + } s := &SessionManager{ + sem: sem, sessions: make(map[string]*builder), sessionTimers: make(map[string]*time.Timer), blockchain: blockchain, @@ -61,12 +73,17 @@ func NewSessionManager(blockchain blockchain, config *Config) *SessionManager { } // NewSession creates a new builder session and returns the session id -func (s *SessionManager) NewSession() (string, error) { - s.sessionsLock.Lock() - defer s.sessionsLock.Unlock() +func (s *SessionManager) NewSession(ctx context.Context) (string, error) { + // Wait for session to become available + select { + case <-s.sem: + s.sessionsLock.Lock() + defer s.sessionsLock.Unlock() + case <-ctx.Done(): + return "", ctx.Err() + } parent := s.blockchain.CurrentHeader() - chainConfig := s.blockchain.Config() header := &types.Header{ @@ -111,6 +128,14 @@ func (s *SessionManager) NewSession() (string, error) { delete(s.sessionTimers, id) }) + // Technically, we are certain that there is an open slot in the semaphore + // channel, but let's be defensive and panic if the invariant is violated. + select { + case s.sem <- struct{}{}: + default: + panic("released more sessions than are open") // unreachable + } + return id, nil } diff --git a/suave/builder/session_manager_test.go b/suave/builder/session_manager_test.go index cdd20a1b43..d09d9dfc88 100644 --- a/suave/builder/session_manager_test.go +++ b/suave/builder/session_manager_test.go @@ -1,6 +1,7 @@ package builder import ( + "context" "crypto/ecdsa" "math/big" "testing" @@ -21,7 +22,7 @@ func TestSessionManager_SessionTimeout(t *testing.T) { SessionIdleTimeout: 500 * time.Millisecond, }) - id, err := mngr.NewSession() + id, err := mngr.NewSession(context.TODO()) require.NoError(t, err) time.Sleep(1 * time.Second) @@ -30,12 +31,47 @@ func TestSessionManager_SessionTimeout(t *testing.T) { require.Error(t, err) } +func TestSessionManager_MaxConcurrentSessions(t *testing.T) { + t.Parallel() + + const d = time.Millisecond * 10 + + mngr, _ := newSessionManager(t, &Config{ + MaxConcurrentSessions: 1, + SessionIdleTimeout: d, + }) + + t.Run("SessionAvailable", func(t *testing.T) { + sess, err := mngr.NewSession(context.TODO()) + require.NoError(t, err) + require.NotZero(t, sess) + }) + + t.Run("ContextExpired", func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + sess, err := mngr.NewSession(ctx) + require.Zero(t, sess) + require.ErrorIs(t, err, context.Canceled) + }) + + t.Run("SessionExpired", func(t *testing.T) { + time.Sleep(d) // Wait for the session to expire. + + // We should be able to open a session again. + sess, err := mngr.NewSession(context.TODO()) + require.NoError(t, err) + require.NotZero(t, sess) + }) +} + func TestSessionManager_SessionRefresh(t *testing.T) { mngr, _ := newSessionManager(t, &Config{ SessionIdleTimeout: 500 * time.Millisecond, }) - id, err := mngr.NewSession() + id, err := mngr.NewSession(context.TODO()) require.NoError(t, err) // if we query the session under the idle timeout, @@ -60,7 +96,7 @@ func TestSessionManager_StartSession(t *testing.T) { // test that the session starts and it can simulate transactions mngr, bMock := newSessionManager(t, &Config{}) - id, err := mngr.NewSession() + id, err := mngr.NewSession(context.TODO()) require.NoError(t, err) txn := bMock.state.newTransfer(t, common.Address{}, big.NewInt(1))