✨ supervisor Add option to run supervisor a fixed number of times (#620)

<!--
Copyright (C) 2020-2022 Arm Limited or its affiliates and Contributors.
All rights reserved.
SPDX-License-Identifier: Apache-2.0
-->
### Description

<!--
Please add any detail or context that would be useful to a reviewer.
-->

Add option to run supervisor a fixed number of times

### Test Coverage

<!--
Please put an `x` in the correct box e.g. `[x]` to indicate the testing
coverage of this change.
-->

- [x]  This change is covered by existing or additional automated tests.
- [ ] Manual testing has been performed (and evidence provided) as
automated testing was not feasible.
- [ ] Additional tests are not required for this change (e.g.
documentation update).

Author: joshjennings98

changes/20250529160336.feature

@@ -0,0 +1 @@
+:sparkles: `supervisor` Add option to run supervisor a fixed number of times

changes/20250529161638.feature

@@ -0,0 +1 @@
+:sparkles: `supervisor` Add option to run a function after the supervisor stops

changes/20250529175329.feature

@@ -0,0 +1 @@
+:sparkles: `logs` Add support for a FIFO logger that discards logs ater reading them

utils/logs/fifo_logger.go

@@ -0,0 +1,91 @@
+package logs
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"log"
+	"sync"
+)
+
+type FIFOWriter struct {
+	io.WriteCloser
+	mu   sync.RWMutex
+	Logs bytes.Buffer
+}
+
+func (w *FIFOWriter) Write(p []byte) (n int, err error) {
+	w.mu.RLock()
+	defer w.mu.RUnlock()
+	w.Logs.Write(p)
+	return
+}
+
+func (w *FIFOWriter) Close() (err error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	w.Logs.Reset()
+	return
+}
+
+func (w *FIFOWriter) Read() string {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	n := w.Logs.Len()
+	if n == 0 {
+		return ""
+	}
+	bytes := w.Logs.Next(n)
+	return string(bytes)
+}
+
+type FIFOLoggers struct {
+	GenericLoggers
+	LogWriter FIFOWriter
+}
+
+func (l *FIFOLoggers) Check() error {
+	return l.GenericLoggers.Check()
+}
+
+func (l *FIFOLoggers) Read() string {
+	return l.LogWriter.Read()
+}
+
+// Close closes the logger
+func (l *FIFOLoggers) Close() (err error) {
+	err = l.LogWriter.Close()
+	if err != nil {
+		return
+	}
+	err = l.GenericLoggers.Close()
+	return
+}
+
+// NewFIFOLogger creates a logger to a bytes buffer.
+// All messages (whether they are output or error) are merged together.
+// Once messages have been accessed they are gone
+func NewFIFOLogger(loggerSource string) (loggers *FIFOLoggers, err error) {
+	loggers = &FIFOLoggers{
+		LogWriter: FIFOWriter{},
+	}
+	loggers.GenericLoggers = GenericLoggers{
+		Output: log.New(&loggers.LogWriter, fmt.Sprintf("[%v] Output: ", loggerSource), log.LstdFlags),
+		Error:  log.New(&loggers.LogWriter, fmt.Sprintf("[%v] Error: ", loggerSource), log.LstdFlags),
+	}
+	return
+}
+
+// NewPlainFIFOLogger creates a logger to a bytes buffer with no extra flag, prefix or tag, just the logged text.
+// All messages (whether they are output or error) are merged together.
+// Once messages have been accessed they are gone
+func NewPlainFIFOLogger() (loggers *FIFOLoggers, err error) {
+	loggers = &FIFOLoggers{
+		LogWriter: FIFOWriter{},
+	}
+	loggers.GenericLoggers = GenericLoggers{
+		Output: log.New(&loggers.LogWriter, "", 0),
+		Error:  log.New(&loggers.LogWriter, "", 0),
+	}
+	return
+}

utils/logs/fifo_logger_test.go

@@ -0,0 +1,48 @@
+package logs
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestFIFOLogger(t *testing.T) {
+	loggers, err := NewFIFOLogger("Test")
+	require.NoError(t, err)
+	testLog(t, loggers)
+	loggers.LogError("Test err")
+	loggers.Log("Test1")
+	contents := loggers.Read()
+	require.NotEmpty(t, contents)
+	require.True(t, strings.Contains(contents, "Test err"))
+	require.True(t, strings.Contains(contents, "Test1"))
+	loggers.Log("Test2")
+	contents = loggers.Read()
+	require.NotEmpty(t, contents)
+	require.False(t, strings.Contains(contents, "Test err"))
+	require.False(t, strings.Contains(contents, "Test1"))
+	require.True(t, strings.Contains(contents, "Test2"))
+	contents = loggers.Read()
+	require.Empty(t, contents)
+}
+
+func TestPlainFIFOLogger(t *testing.T) {
+	loggers, err := NewPlainFIFOLogger()
+	require.NoError(t, err)
+	testLog(t, loggers)
+	loggers.LogError("Test err")
+	loggers.Log("Test1")
+	contents := loggers.Read()
+	require.NotEmpty(t, contents)
+	require.True(t, strings.Contains(contents, "Test err"))
+	require.True(t, strings.Contains(contents, "Test1"))
+	loggers.Log("Test2")
+	contents = loggers.Read()
+	require.NotEmpty(t, contents)
+	require.False(t, strings.Contains(contents, "Test err"))
+	require.False(t, strings.Contains(contents, "Test1"))
+	require.True(t, strings.Contains(contents, "Test2"))
+	contents = loggers.Read()
+	require.Empty(t, contents)
+}

utils/subprocess/supervisor/supervisor.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/ARM-software/golang-utils/utils/commonerrors"
 	"github.com/ARM-software/golang-utils/utils/parallelisation"
+	"github.com/ARM-software/golang-utils/utils/safecast"
 	"github.com/ARM-software/golang-utils/utils/subprocess"
 )
 
@@ -19,8 +20,10 @@ type Supervisor struct {
 	preStart      func(context.Context) error
 	postStart     func(context.Context) error
 	postStop      func(context.Context, error) error
+	postEnd       func()
 	haltingErrors []error
 	restartDelay  time.Duration
+	count         uint
 }
 
 type SupervisorOption func(*Supervisor)
@@ -75,9 +78,29 @@ func WithRestartDelay(delay time.Duration) SupervisorOption {
 	}
 }
 
+// WithCount will run cause the supervisor to exit after 'count' executions.
+func WithCount[I safecast.INumber](count I) SupervisorOption {
+	return func(s *Supervisor) {
+		s.count = safecast.ToUint(count)
+	}
+}
+
+// WithPostEnd will run 'function' after the supervisor has stopped.
+// It does not take a context to ensure that it runs after a context has been cancelled.
+// It does not return an error as this could cause confusion with the other returned errors.
+func WithPostEnd(function func()) SupervisorOption {
+	return func(s *Supervisor) {
+		s.postEnd = function
+	}
+}
+
 // Run will run the supervisor and execute any of the command hooks. If it receives a halting error or the context is cancelled then it will exit
 func (s *Supervisor) Run(ctx context.Context) (err error) {
-	for {
+	if s.postEnd != nil {
+		defer s.postEnd()
+	}
+
+	for i := uint(0); s.count == 0 || i < s.count; i++ {
 		err = parallelisation.DetermineContextError(ctx)
 		if err != nil {
 			return
@@ -141,4 +164,6 @@ func (s *Supervisor) Run(ctx context.Context) (err error) {
 
 		// restart process
 	}
+
+	return
 }

utils/subprocess/supervisor/supervisor_test.go

@@ -378,4 +378,61 @@ func TestSupervisor(t *testing.T) {
 		assert.NotEmpty(t, written)
 		assert.Equal(t, string(written), "test\ntest123\n")
 	})
+
+	t.Run("with count", func(t *testing.T) {
+		if platform.IsWindows() {
+			t.SkipNow()
+		}
+
+		ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+		defer cancel()
+
+		logger, err := logs.NewLogrLogger(logstest.NewTestLogger(t), "Test")
+		require.NoError(t, err)
+
+		counter := atomic.Uint64{}
+		assert.Zero(t, counter.Load())
+
+		cmd, err := subprocess.New(ctx, logger, "starting", "success", "failed", "echo", "123")
+		require.NoError(t, err)
+
+		runner := NewSupervisor(func(ctx context.Context) (*subprocess.Subprocess, error) {
+			return cmd, nil
+		}, WithPostStop(func(_ context.Context, _ error) error {
+			_ = counter.Inc()
+			return nil
+		}), WithCount(3))
+
+		err = runner.Run(ctx)
+		assert.NoError(t, err)
+		assert.EqualValues(t, 3, counter.Load())
+	})
+
+	t.Run("with post end", func(t *testing.T) {
+		if platform.IsWindows() {
+			t.SkipNow()
+		}
+
+		ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+		defer cancel()
+
+		logger, err := logs.NewLogrLogger(logstest.NewTestLogger(t), "Test")
+		require.NoError(t, err)
+
+		counter := atomic.Uint64{}
+		assert.Zero(t, counter.Load())
+
+		cmd, err := subprocess.New(ctx, logger, "starting", "success", "failed", "echo", "123")
+		require.NoError(t, err)
+
+		runner := NewSupervisor(func(ctx context.Context) (*subprocess.Subprocess, error) {
+			return cmd, nil
+		}, WithPostEnd(func() {
+			_ = counter.Inc()
+		}), WithCount(1))
+
+		err = runner.Run(ctx)
+		assert.NoError(t, err)
+		assert.EqualValues(t, 1, counter.Load())
+	})
 }