[#72940] linux-client: Add ShellRunner

ShellRunner tracks all shell sessions managed by the daemon.

Signed-off-by: Łukasz Kędziora <lkedziora@antmicro.com>

Author: lkedziora

devices/linux-client/shell/runner.go

@@ -0,0 +1,65 @@
+package shell
+
+import (
+	"errors"
+	"sync"
+
+	"golang.org/x/sync/semaphore"
+)
+
+// Keeps track of shell processes spawned by the daemon.
+type ShellRunner struct {
+	shellCount *semaphore.Weighted      // Current available shell count.
+	sessions   map[string]*ShellSession // Dict of all currently running shells.
+	mut        sync.Mutex               // Protects sessions dict.
+}
+
+var (
+	ErrSessionNotFound     = errors.New("session not found")
+	ErrSessionLimitReached = errors.New("session limit reached")
+)
+
+// Create a new shell runner that is capable of spawning up to limit concurrent
+// shell processes.
+func NewShellRunner(limit int) (*ShellRunner, error) {
+	sr := new(ShellRunner)
+	sr.shellCount = semaphore.NewWeighted(int64(limit))
+	sr.sessions = make(map[string]*ShellSession)
+	return sr, nil
+}
+
+// Spawn a new shell process and associate it with a given UUID.
+func (s *ShellRunner) Spawn(uuid string) (*ShellSession, error) {
+	ok := s.shellCount.TryAcquire(1)
+	if !ok {
+		return nil, ErrSessionLimitReached
+	}
+
+	session, err := NewShellSession(uuid)
+	if err != nil {
+		return nil, err
+	}
+
+	s.mut.Lock()
+	defer s.mut.Unlock()
+	s.sessions[uuid] = session
+
+	return session, nil
+}
+
+// Terminate a shell session that was previously spawned. The associated shell
+// process is terminated if it is still running.
+func (s *ShellRunner) Terminate(uuid string) error {
+	s.mut.Lock()
+	defer s.mut.Unlock()
+
+	v, ok := s.sessions[uuid]
+	if !ok {
+		return ErrSessionNotFound
+	}
+	v.Close()
+	delete(s.sessions, uuid)
+	s.shellCount.Release(1)
+
+	return nil
+}

devices/linux-client/shell/runner_test.go

@@ -0,0 +1,62 @@
+package shell
+
+import (
+	"strconv"
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+)
+
+const (
+	TestSessionUuid = "ABCDEF"
+)
+
+func TestShellRunnerBasic(t *testing.T) {
+	sr, err := NewShellRunner(1)
+	assert.NoError(t, err)
+
+	_, err = sr.Spawn(TestSessionUuid)
+	assert.NoError(t, err)
+}
+
+func TestShellRunnerLimits(t *testing.T) {
+	const (
+		LimitToTest = 4
+	)
+
+	sr, err := NewShellRunner(LimitToTest)
+	assert.NoError(t, err)
+
+	for i := 0; i < LimitToTest; i += 1 {
+		_, err = sr.Spawn(strconv.Itoa(i))
+		assert.NoError(t, err)
+	}
+
+	_, err = sr.Spawn(uuid.NewString())
+	if assert.Error(t, err) {
+		assert.Equal(t, ErrSessionLimitReached, err)
+	}
+}
+
+func TestShellRunnerLimitsIncrementedAfterTermination(t *testing.T) {
+	sr, err := NewShellRunner(1)
+	assert.NoError(t, err)
+
+	id1 := uuid.NewString()
+	id2 := uuid.NewString()
+
+	_, err = sr.Spawn(id1)
+	assert.NoError(t, err)
+
+	_, err = sr.Spawn(id2)
+	if assert.Error(t, err) {
+		assert.Equal(t, ErrSessionLimitReached, err)
+	}
+
+	err = sr.Terminate(id1)
+	assert.NoError(t, err)
+
+	_, err = sr.Spawn(id2)
+	assert.NoError(t, err)
+}