Fix handling sessions in runnerv2 service (#579)

Closes #567

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
+}

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])
 }

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)
+}

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)
+}