From accac4912aa8dc303d56988677037346d79f279f Mon Sep 17 00:00:00 2001 From: Adam Babik Date: Fri, 17 May 2024 21:36:24 +0200 Subject: [PATCH] Fix handling sessions in runnerv2 service (#579) Closes https://github.com/stateful/runme/issues/567 --- internal/command/command_test.go | 20 +++++++- internal/command/session.go | 46 +++-------------- internal/command/session_test.go | 33 ++++++++++++ .../runnerv2service/service_sessions_test.go | 50 +++++++++++++++++++ 4 files changed, 107 insertions(+), 42 deletions(-) create mode 100644 internal/command/session_test.go diff --git a/internal/command/command_test.go b/internal/command/command_test.go index f8bdb8dd9..00032c07c 100644 --- a/internal/command/command_test.go +++ b/internal/command/command_test.go @@ -195,7 +195,7 @@ func testExecuteNativeCommand( options := Options{ Logger: logger, - Session: MustNewSessionWithEnv(env...), + Session: mustNewSessionWithEnv(env...), Stdout: stdout, Stderr: stderr, Stdin: input, @@ -237,7 +237,7 @@ func testExecuteVirtualCommand( stdout := bytes.NewBuffer(nil) options := Options{ - Session: MustNewSessionWithEnv(env...), + Session: mustNewSessionWithEnv(env...), Stdout: stdout, Stdin: input, Logger: logger, @@ -328,3 +328,19 @@ func TestCommandWithSession(t *testing.T) { require.Equal(t, "TEST_ENV equals test1", stdout.String()) }) } + +func newSessionWithEnv(env ...string) (*Session, error) { + s := NewSession() + if err := s.SetEnv(env...); err != nil { + return nil, err + } + return s, nil +} + +func mustNewSessionWithEnv(env ...string) *Session { + s, err := newSessionWithEnv(env...) + if err != nil { + panic(err) + } + return s +} diff --git a/internal/command/session.go b/internal/command/session.go index 583944a96..08eeb417d 100644 --- a/internal/command/session.go +++ b/internal/command/session.go @@ -1,8 +1,6 @@ package command import ( - "sync" - lru "github.com/hashicorp/golang-lru/v2" "github.com/stateful/runme/v3/internal/ulid" @@ -23,22 +21,6 @@ func NewSession() *Session { } } -func NewSessionWithEnv(env ...string) (*Session, error) { - s := NewSession() - if err := s.SetEnv(env...); err != nil { - return nil, err - } - return s, nil -} - -func MustNewSessionWithEnv(env ...string) *Session { - s, err := NewSessionWithEnv(env...) - if err != nil { - panic(err) - } - return s -} - func (s *Session) SetEnv(env ...string) error { _, err := s.envStore.Merge(env...) return err @@ -63,10 +45,6 @@ const SessionListCapacity = 1024 type SessionList struct { items *lru.Cache[string, *Session] - // Even though, the lry.Cache is thread-safe, it does not provide a way to - // get the most recent added session. - mu sync.Mutex - latestSession *Session } func NewSessionList() (*SessionList, error) { @@ -79,10 +57,6 @@ func NewSessionList() (*SessionList, error) { func (sl *SessionList) Add(session *Session) { sl.items.Add(session.ID, session) - - sl.mu.Lock() - sl.latestSession = session - sl.mu.Unlock() } func (sl *SessionList) Get(id string) (*Session, bool) { @@ -94,21 +68,13 @@ func (sl *SessionList) List() []*Session { } func (sl *SessionList) Delete(id string) bool { - ok := sl.items.Remove(id) - if !ok { - return ok - } - - sl.mu.Lock() - if sl.latestSession != nil && sl.latestSession.ID == id { - keys := sl.items.Keys() - sl.latestSession, _ = sl.items.Get(keys[len(keys)-1]) - } - sl.mu.Unlock() - - return ok + return sl.items.Remove(id) } func (sl *SessionList) Newest() (*Session, bool) { - return sl.latestSession, sl.latestSession != nil + keys := sl.items.Keys() + if len(keys) == 0 { + return nil, false + } + return sl.items.Get(keys[len(keys)-1]) } diff --git a/internal/command/session_test.go b/internal/command/session_test.go new file mode 100644 index 000000000..f9bc6489b --- /dev/null +++ b/internal/command/session_test.go @@ -0,0 +1,33 @@ +package command + +import ( + "testing" + + "github.com/stretchr/testify/require" + "golang.org/x/sync/errgroup" +) + +func TestSessionList(t *testing.T) { + l, err := NewSessionList() + require.NoError(t, err) + + var g errgroup.Group + + for i := 0; i < 10; i++ { + g.Go(func() error { + s := NewSession() + l.Add(s) + return nil + }) + } + + require.NoError(t, g.Wait()) + require.Len(t, l.items.Keys(), 10) + + s := NewSession() + l.Add(s) + + newest, ok := l.Newest() + require.True(t, ok) + require.Equal(t, s.ID, newest.ID) +} diff --git a/internal/runnerv2service/service_sessions_test.go b/internal/runnerv2service/service_sessions_test.go index a0cea8565..41c699170 100644 --- a/internal/runnerv2service/service_sessions_test.go +++ b/internal/runnerv2service/service_sessions_test.go @@ -6,6 +6,7 @@ import ( "context" "testing" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/stateful/runme/v3/internal/project/testdata" @@ -66,3 +67,52 @@ func TestRunnerServiceSessions(t *testing.T) { require.Error(t, err) }) } + +func TestRunnerServiceSessions_StrategyMostRecent(t *testing.T) { + lis, stop := testStartRunnerServiceServer(t) + t.Cleanup(stop) + _, client := testCreateRunnerServiceClient(t, lis) + + // Create a session with env. + sessResp, err := client.CreateSession( + context.Background(), + &runnerv2alpha1.CreateSessionRequest{ + Env: []string{"TEST1=value1"}, + }, + ) + require.NoError(t, err) + + // Prep the execute stream. + stream, err := client.Execute(context.Background()) + require.NoError(t, err) + + execResult := make(chan executeResult) + go getExecuteResult(stream, execResult) + + // Execute a program using the most recent session strategy. + req := &runnerv2alpha1.ExecuteRequest{ + Config: &runnerv2alpha1.ProgramConfig{ + ProgramName: "bash", + Source: &runnerv2alpha1.ProgramConfig_Commands{ + Commands: &runnerv2alpha1.ProgramConfig_CommandList{ + Items: []string{ + `echo "TEST1=$TEST1"`, + }, + }, + }, + }, + SessionStrategy: runnerv2alpha1.SessionStrategy_SESSION_STRATEGY_MOST_RECENT, + } + err = stream.Send(req) + require.NoError(t, err) + + result := <-execResult + + assert.NoError(t, result.Err) + assert.EqualValues(t, 0, result.ExitCode) + assert.Equal(t, "TEST1=value1\n", string(result.Stdout)) + + // Delete the session. + _, err = client.DeleteSession(context.Background(), &runnerv2alpha1.DeleteSessionRequest{Id: sessResp.GetSession().GetId()}) + require.NoError(t, err) +}