Add additional performance spans

internal/api/environments.go

@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"github.com/getsentry/sentry-go"
 	"github.com/gorilla/mux"
 	"github.com/openHPI/poseidon/internal/environment"
 	"github.com/openHPI/poseidon/internal/runner"
@@ -118,7 +119,9 @@ func (e *EnvironmentController) createOrUpdate(writer http.ResponseWriter, reque
 		return
 	}
 
-	created, err := e.manager.CreateOrUpdate(environmentID, *req)
+	span := sentry.StartSpan(request.Context(), "Create Environment")
+	created, err := e.manager.CreateOrUpdate(environmentID, *req, request.Context())
+	span.Finish()
 	if err != nil {
 		writeInternalServerError(writer, err, dto.ErrorUnknown)
 	}

internal/api/environments_test.go

@@ -246,7 +246,7 @@ func (s *CreateOrUpdateEnvironmentTestSuite) TestReturnsBadRequestWhenBadBody()
 func (s *CreateOrUpdateEnvironmentTestSuite) TestReturnsInternalServerErrorWhenManagerReturnsError() {
 	testError := tests.ErrDefault
 	s.manager.
-		On("CreateOrUpdate", s.id, mock.AnythingOfType("dto.ExecutionEnvironmentRequest")).
+		On("CreateOrUpdate", s.id, mock.AnythingOfType("dto.ExecutionEnvironmentRequest"), mock.Anything).
 		Return(false, testError)
 
 	recorder := s.recordRequest()
@@ -256,7 +256,7 @@ func (s *CreateOrUpdateEnvironmentTestSuite) TestReturnsInternalServerErrorWhenM
 
 func (s *CreateOrUpdateEnvironmentTestSuite) TestReturnsCreatedIfNewEnvironment() {
 	s.manager.
-		On("CreateOrUpdate", s.id, mock.AnythingOfType("dto.ExecutionEnvironmentRequest")).
+		On("CreateOrUpdate", s.id, mock.AnythingOfType("dto.ExecutionEnvironmentRequest"), mock.Anything).
 		Return(true, nil)
 
 	recorder := s.recordRequest()
@@ -265,7 +265,7 @@ func (s *CreateOrUpdateEnvironmentTestSuite) TestReturnsCreatedIfNewEnvironment(
 
 func (s *CreateOrUpdateEnvironmentTestSuite) TestReturnsNoContentIfNotNewEnvironment() {
 	s.manager.
-		On("CreateOrUpdate", s.id, mock.AnythingOfType("dto.ExecutionEnvironmentRequest")).
+		On("CreateOrUpdate", s.id, mock.AnythingOfType("dto.ExecutionEnvironmentRequest"), mock.Anything).
 		Return(false, nil)
 
 	recorder := s.recordRequest()

internal/api/runners.go

@@ -3,6 +3,7 @@ package api
 import (
 	"errors"
 	"fmt"
+	"github.com/getsentry/sentry-go"
 	"github.com/google/uuid"
 	"github.com/gorilla/mux"
 	"github.com/openHPI/poseidon/internal/config"
@@ -65,7 +66,9 @@ func (r *RunnerController) provide(writer http.ResponseWriter, request *http.Req
 	}
 	environmentID := dto.EnvironmentID(runnerRequest.ExecutionEnvironmentID)
 
+	span := sentry.StartSpan(request.Context(), "Claim Runner")
 	nextRunner, err := r.manager.Claim(environmentID, runnerRequest.InactivityTimeout)
+	span.Finish()
 	if err != nil {
 		switch {
 		case errors.Is(err, runner.ErrUnknownExecutionEnvironment):
@@ -103,7 +106,9 @@ func (r *RunnerController) listFileSystem(writer http.ResponseWriter, request *h
 	}
 
 	writer.Header().Set("Content-Type", "application/json")
+	span := sentry.StartSpan(request.Context(), "List File System")
 	err = targetRunner.ListFileSystem(path, recursive, writer, privilegedExecution, request.Context())
+	span.Finish()
 	if errors.Is(err, runner.ErrFileNotFound) {
 		writeClientError(writer, err, http.StatusFailedDependency)
 		return
@@ -125,7 +130,11 @@ func (r *RunnerController) updateFileSystem(writer http.ResponseWriter, request
 
 	targetRunner, _ := runner.FromContext(request.Context())
 	monitoring.AddRunnerMonitoringData(request, targetRunner.ID(), targetRunner.Environment())
-	if err := targetRunner.UpdateFileSystem(fileCopyRequest); err != nil {
+
+	span := sentry.StartSpan(request.Context(), "Update File System")
+	err := targetRunner.UpdateFileSystem(fileCopyRequest, request.Context())
+	span.Finish()
+	if err != nil {
 		log.WithError(err).Error("Could not perform the requested updateFileSystem.")
 		writeInternalServerError(writer, err, dto.ErrorUnknown)
 		return
@@ -144,7 +153,9 @@ func (r *RunnerController) fileContent(writer http.ResponseWriter, request *http
 	}
 
 	writer.Header().Set("Content-Disposition", "attachment; filename=\""+path+"\"")
+	span := sentry.StartSpan(request.Context(), "File Content")
 	err = targetRunner.GetFileContent(path, writer, privilegedExecution, request.Context())
+	span.Finish()
 	if errors.Is(err, runner.ErrFileNotFound) {
 		writeClientError(writer, err, http.StatusFailedDependency)
 		return
@@ -191,7 +202,9 @@ func (r *RunnerController) execute(writer http.ResponseWriter, request *http.Req
 		return
 	}
 	id := newUUID.String()
+	span := sentry.StartSpan(request.Context(), "Store Execution")
 	targetRunner.StoreExecution(id, executionRequest)
+	span.Finish()
 	webSocketURL := url.URL{
 		Scheme:   scheme,
 		Host:     request.Host,
@@ -230,7 +243,9 @@ func (r *RunnerController) delete(writer http.ResponseWriter, request *http.Requ
 	targetRunner, _ := runner.FromContext(request.Context())
 	monitoring.AddRunnerMonitoringData(request, targetRunner.ID(), targetRunner.Environment())
 
+	span := sentry.StartSpan(request.Context(), "Return Runner")
 	err := r.manager.Return(targetRunner)
+	span.Finish()
 	if err != nil {
 		writeInternalServerError(writer, err, dto.ErrorNomadInternalServerError)
 		return

internal/api/runners_test.go

@@ -277,7 +277,8 @@ func (s *UpdateFileSystemRouteTestSuite) SetupTest() {
 }
 
 func (s *UpdateFileSystemRouteTestSuite) TestUpdateFileSystemReturnsNoContentOnValidRequest() {
-	s.runnerMock.On("UpdateFileSystem", mock.AnythingOfType("*dto.UpdateFileSystemRequest")).Return(nil)
+	s.runnerMock.On("UpdateFileSystem", mock.AnythingOfType("*dto.UpdateFileSystemRequest"), mock.Anything).
+		Return(nil)
 
 	copyRequest := dto.UpdateFileSystemRequest{}
 	body, err := json.Marshal(copyRequest)
@@ -287,7 +288,8 @@ func (s *UpdateFileSystemRouteTestSuite) TestUpdateFileSystemReturnsNoContentOnV
 
 	s.router.ServeHTTP(s.recorder, request)
 	s.Equal(http.StatusNoContent, s.recorder.Code)
-	s.runnerMock.AssertCalled(s.T(), "UpdateFileSystem", mock.AnythingOfType("*dto.UpdateFileSystemRequest"))
+	s.runnerMock.AssertCalled(s.T(), "UpdateFileSystem",
+		mock.AnythingOfType("*dto.UpdateFileSystemRequest"), mock.Anything)
 }
 
 func (s *UpdateFileSystemRouteTestSuite) TestUpdateFileSystemReturnsBadRequestOnInvalidRequestBody() {
@@ -314,7 +316,7 @@ func (s *UpdateFileSystemRouteTestSuite) TestUpdateFileSystemToNonExistingRunner
 
 func (s *UpdateFileSystemRouteTestSuite) TestUpdateFileSystemReturnsInternalServerErrorWhenCopyFailed() {
 	s.runnerMock.
-		On("UpdateFileSystem", mock.AnythingOfType("*dto.UpdateFileSystemRequest")).
+		On("UpdateFileSystem", mock.AnythingOfType("*dto.UpdateFileSystemRequest"), mock.Anything).
 		Return(runner.ErrorFileCopyFailed)
 
 	copyRequest := dto.UpdateFileSystemRequest{}

internal/api/websocket.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"github.com/getsentry/sentry-go"
 	"github.com/gorilla/websocket"
 	"github.com/openHPI/poseidon/internal/api/ws"
 	"github.com/openHPI/poseidon/internal/runner"
@@ -96,6 +97,8 @@ func (r *RunnerController) connectToRunner(writer http.ResponseWriter, request *
 	log.WithField("runnerId", targetRunner.ID()).
 		WithField("executionID", logging.RemoveNewlineSymbol(executionID)).
 		Info("Running execution")
+	span := sentry.StartSpan(request.Context(), "Execute Interactively")
+	defer span.Finish()
 	exit, cancel, err := targetRunner.ExecuteInteractively(executionID,
 		proxy.Input, proxy.Output.StdOut(), proxy.Output.StdErr())
 	if err != nil {

internal/environment/abstract_manager.go

@@ -1,6 +1,7 @@
 package environment
 
 import (
+	"context"
 	"fmt"
 	"github.com/openHPI/poseidon/internal/runner"
 	"github.com/openHPI/poseidon/pkg/dto"
@@ -37,7 +38,8 @@ func (n *AbstractManager) Get(_ dto.EnvironmentID, _ bool) (runner.ExecutionEnvi
 	return nil, runner.ErrRunnerNotFound
 }
 
-func (n *AbstractManager) CreateOrUpdate(_ dto.EnvironmentID, _ dto.ExecutionEnvironmentRequest) (bool, error) {
+func (n *AbstractManager) CreateOrUpdate(_ dto.EnvironmentID, _ dto.ExecutionEnvironmentRequest, _ context.Context) (
+	bool, error) {
 	return false, nil
 }
 

internal/environment/aws_manager.go

@@ -1,6 +1,7 @@
 package environment
 
 import (
+	"context"
 	"fmt"
 	"github.com/openHPI/poseidon/internal/config"
 	"github.com/openHPI/poseidon/internal/runner"
@@ -41,9 +42,9 @@ func (a *AWSEnvironmentManager) Get(id dto.EnvironmentID, fetch bool) (runner.Ex
 }
 
 func (a *AWSEnvironmentManager) CreateOrUpdate(
-	id dto.EnvironmentID, request dto.ExecutionEnvironmentRequest) (bool, error) {
+	id dto.EnvironmentID, request dto.ExecutionEnvironmentRequest, ctx context.Context) (bool, error) {
 	if !isAWSEnvironment(request) {
-		isCreated, err := a.NextHandler().CreateOrUpdate(id, request)
+		isCreated, err := a.NextHandler().CreateOrUpdate(id, request, ctx)
 		if err != nil {
 			return false, fmt.Errorf("aws wrapped: %w", err)
 		}

internal/environment/aws_manager_test.go

@@ -1,6 +1,7 @@
 package environment
 
 import (
+	"context"
 	"github.com/openHPI/poseidon/internal/config"
 	"github.com/openHPI/poseidon/internal/runner"
 	"github.com/openHPI/poseidon/pkg/dto"
@@ -18,7 +19,8 @@ func TestAWSEnvironmentManager_CreateOrUpdate(t *testing.T) {
 
 	t.Run("can create default Java environment", func(t *testing.T) {
 		config.Config.AWS.Functions = []string{uniqueImage}
-		_, err := m.CreateOrUpdate(tests.AnotherEnvironmentIDAsInteger, dto.ExecutionEnvironmentRequest{Image: uniqueImage})
+		_, err := m.CreateOrUpdate(
+			tests.AnotherEnvironmentIDAsInteger, dto.ExecutionEnvironmentRequest{Image: uniqueImage}, context.Background())
 		assert.NoError(t, err)
 	})
 
@@ -31,14 +33,14 @@ func TestAWSEnvironmentManager_CreateOrUpdate(t *testing.T) {
 	t.Run("non-handleable requests are forwarded to the next manager", func(t *testing.T) {
 		nextHandler := &ManagerHandlerMock{}
 		nextHandler.On("CreateOrUpdate", mock.AnythingOfType("dto.EnvironmentID"),
-			mock.AnythingOfType("dto.ExecutionEnvironmentRequest")).Return(true, nil)
+			mock.AnythingOfType("dto.ExecutionEnvironmentRequest"), mock.Anything).Return(true, nil)
 		m.SetNextHandler(nextHandler)
 
 		request := dto.ExecutionEnvironmentRequest{}
-		_, err := m.CreateOrUpdate(tests.DefaultEnvironmentIDAsInteger, request)
+		_, err := m.CreateOrUpdate(tests.DefaultEnvironmentIDAsInteger, request, context.Background())
 		assert.NoError(t, err)
 		nextHandler.AssertCalled(t, "CreateOrUpdate",
-			dto.EnvironmentID(tests.DefaultEnvironmentIDAsInteger), request)
+			dto.EnvironmentID(tests.DefaultEnvironmentIDAsInteger), request, mock.Anything)
 	})
 }
 

internal/environment/manager.go

@@ -1,6 +1,7 @@
 package environment
 
 import (
+	"context"
 	"github.com/openHPI/poseidon/internal/runner"
 	"github.com/openHPI/poseidon/pkg/dto"
 )
@@ -30,6 +31,7 @@ type Manager interface {
 	CreateOrUpdate(
 		id dto.EnvironmentID,
 		request dto.ExecutionEnvironmentRequest,
+		ctx context.Context,
 	) (bool, error)
 
 	// Delete removes the specified execution environment.

internal/environment/manager_handler_mock.go

@@ -1,8 +1,10 @@
-// Code generated by mockery v2.12.3. DO NOT EDIT.
+// Code generated by mockery v2.16.0. DO NOT EDIT.
 
 package environment
 
 import (
+	context "context"
+
 	dto "github.com/openHPI/poseidon/pkg/dto"
 	mock "github.com/stretchr/testify/mock"
 
@@ -14,20 +16,20 @@ type ManagerHandlerMock struct {
 	mock.Mock
 }
 
-// CreateOrUpdate provides a mock function with given fields: id, request
+// CreateOrUpdate provides a mock function with given fields: id, request, ctx
-func (_m *ManagerHandlerMock) CreateOrUpdate(id dto.EnvironmentID, request dto.ExecutionEnvironmentRequest) (bool, error) {
+func (_m *ManagerHandlerMock) CreateOrUpdate(id dto.EnvironmentID, request dto.ExecutionEnvironmentRequest, ctx context.Context) (bool, error) {
-	ret := _m.Called(id, request)
+	ret := _m.Called(id, request, ctx)
 
 	var r0 bool
-	if rf, ok := ret.Get(0).(func(dto.EnvironmentID, dto.ExecutionEnvironmentRequest) bool); ok {
+	if rf, ok := ret.Get(0).(func(dto.EnvironmentID, dto.ExecutionEnvironmentRequest, context.Context) bool); ok {
-		r0 = rf(id, request)
+		r0 = rf(id, request, ctx)
 	} else {
 		r0 = ret.Get(0).(bool)
 	}
 
 	var r1 error
-	if rf, ok := ret.Get(1).(func(dto.EnvironmentID, dto.ExecutionEnvironmentRequest) error); ok {
+	if rf, ok := ret.Get(1).(func(dto.EnvironmentID, dto.ExecutionEnvironmentRequest, context.Context) error); ok {
-		r1 = rf(id, request)
+		r1 = rf(id, request, ctx)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -153,13 +155,13 @@ func (_m *ManagerHandlerMock) Statistics() map[dto.EnvironmentID]*dto.Statistica
 	return r0
 }
 
-type NewManagerHandlerMockT interface {
+type mockConstructorTestingTNewManagerHandlerMock interface {
 	mock.TestingT
 	Cleanup(func())
 }
 
 // NewManagerHandlerMock creates a new instance of ManagerHandlerMock. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
-func NewManagerHandlerMock(t NewManagerHandlerMockT) *ManagerHandlerMock {
+func NewManagerHandlerMock(t mockConstructorTestingTNewManagerHandlerMock) *ManagerHandlerMock {
 	mock := &ManagerHandlerMock{}
 	mock.Mock.Test(t)
 

internal/environment/nomad_manager.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	_ "embed"
 	"fmt"
+	"github.com/getsentry/sentry-go"
 	nomadApi "github.com/hashicorp/nomad/api"
 	"github.com/hashicorp/nomad/nomad/structs"
 	"github.com/openHPI/poseidon/internal/nomad"
@@ -89,8 +90,8 @@ func (m *NomadEnvironmentManager) List(fetch bool) ([]runner.ExecutionEnvironmen
 	return m.runnerManager.ListEnvironments(), nil
 }
 
-func (m *NomadEnvironmentManager) CreateOrUpdate(id dto.EnvironmentID, request dto.ExecutionEnvironmentRequest) (
+func (m *NomadEnvironmentManager) CreateOrUpdate(
-	created bool, err error) {
+	id dto.EnvironmentID, request dto.ExecutionEnvironmentRequest, ctx context.Context) (created bool, err error) {
 	// Check if execution environment is already existing (in the local memory).
 	environment, isExistingEnvironment := m.runnerManager.GetEnvironment(id)
 	if isExistingEnvironment {
@@ -112,13 +113,17 @@ func (m *NomadEnvironmentManager) CreateOrUpdate(id dto.EnvironmentID, request d
 	m.runnerManager.StoreEnvironment(environment)
 
 	// Register template Job with Nomad.
+	span := sentry.StartSpan(ctx, "Register Environment")
 	err = environment.Register()
+	span.Finish()
 	if err != nil {
 		return false, fmt.Errorf("error registering template job in API: %w", err)
 	}
 
 	// Launch idle runners based on the template job.
+	span = sentry.StartSpan(ctx, "Apply Prewarming Pool Size")
 	err = environment.ApplyPrewarmingPoolSize()
+	span.Finish()
 	if err != nil {
 		return false, fmt.Errorf("error scaling template job in API: %w", err)
 	}

internal/environment/nomad_manager_test.go

@@ -59,7 +59,8 @@ func (s *CreateOrUpdateTestSuite) TestReturnsErrorIfCreatesOrUpdateEnvironmentRe
 	s.apiMock.On("DeleteJob", mock.AnythingOfType("string")).Return(nil)
 	s.runnerManagerMock.On("GetEnvironment", mock.AnythingOfType("dto.EnvironmentID")).Return(nil, false)
 	s.runnerManagerMock.On("StoreEnvironment", mock.AnythingOfType("*environment.NomadEnvironment")).Return(true)
-	_, err := s.manager.CreateOrUpdate(dto.EnvironmentID(tests.DefaultEnvironmentIDAsInteger), s.request)
+	_, err := s.manager.CreateOrUpdate(
+		dto.EnvironmentID(tests.DefaultEnvironmentIDAsInteger), s.request, context.Background())
 	s.ErrorIs(err, tests.ErrDefault)
 }
 
@@ -87,7 +88,8 @@ func (s *CreateOrUpdateTestSuite) TestCreateOrUpdatesSetsForcePullFlag() {
 
 		call.ReturnArguments = mock.Arguments{nil}
 	})
-	_, err := s.manager.CreateOrUpdate(dto.EnvironmentID(tests.DefaultEnvironmentIDAsInteger), s.request)
+	_, err := s.manager.CreateOrUpdate(
+		dto.EnvironmentID(tests.DefaultEnvironmentIDAsInteger), s.request, context.Background())
 	s.NoError(err)
 	s.True(count > 1)
 }

internal/runner/aws_runner.go

@@ -117,7 +117,7 @@ func (w *AWSFunctionWorkload) ListFileSystem(_ string, _ bool, _ io.Writer, _ bo
 // Current limitation: No files can be deleted apart from the previously added files.
 // Future Work: Deduplication of the file systems, as the largest workload is likely to be used by additional
 // CSV files or similar, which are the same for many executions.
-func (w *AWSFunctionWorkload) UpdateFileSystem(request *dto.UpdateFileSystemRequest) error {
+func (w *AWSFunctionWorkload) UpdateFileSystem(request *dto.UpdateFileSystemRequest, _ context.Context) error {
 	for _, path := range request.Delete {
 		delete(w.fs, path)
 	}

internal/runner/aws_runner_test.go

@@ -121,7 +121,7 @@ func TestAWSFunctionWorkload_UpdateFileSystem(t *testing.T) {
 	r.StoreExecution(tests.DefaultEnvironmentIDAsString, request)
 	myFile := dto.File{Path: "myPath", Content: []byte("myContent")}
 
-	err = r.UpdateFileSystem(&dto.UpdateFileSystemRequest{Copy: []dto.File{myFile}})
+	err = r.UpdateFileSystem(&dto.UpdateFileSystemRequest{Copy: []dto.File{myFile}}, context.Background())
 	assert.NoError(t, err)
 	_, execCancel, err := r.ExecuteInteractively(tests.DefaultEnvironmentIDAsString, nil, io.Discard, io.Discard)
 	require.NoError(t, err)

internal/runner/nomad_runner.go

@@ -8,6 +8,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"github.com/getsentry/sentry-go"
 	nomadApi "github.com/hashicorp/nomad/api"
 	influxdb2 "github.com/influxdata/influxdb-client-go/v2"
 	"github.com/openHPI/poseidon/internal/nomad"
@@ -153,7 +154,7 @@ func (r *NomadJob) ListFileSystem(
 	return err
 }
 
-func (r *NomadJob) UpdateFileSystem(copyRequest *dto.UpdateFileSystemRequest) error {
+func (r *NomadJob) UpdateFileSystem(copyRequest *dto.UpdateFileSystemRequest, ctx context.Context) error {
 	r.ResetTimeout()
 
 	var tarBuffer bytes.Buffer
@@ -166,10 +167,11 @@ func (r *NomadJob) UpdateFileSystem(copyRequest *dto.UpdateFileSystemRequest) er
 	updateFileCommand := (&dto.ExecutionRequest{Command: fileDeletionCommand + copyCommand}).FullCommand()
 	stdOut := bytes.Buffer{}
 	stdErr := bytes.Buffer{}
+	span := sentry.StartSpan(ctx, "Execute Update File System")
 	exitCode, err := r.api.ExecuteCommand(r.id, context.Background(), updateFileCommand, false,
 		nomad.PrivilegedExecution, // All files should be written and owned by a privileged user #211.
 		&tarBuffer, &stdOut, &stdErr)
-
+	span.Finish()
 	if err != nil {
 		return fmt.Errorf(
 			"%w: nomad error during file copy: %v",

internal/runner/nomad_runner_test.go

@@ -274,7 +274,7 @@ func (s *UpdateFileSystemTestSuite) TestUpdateFileSystemForRunnerPerformsTarExtr
 	// note: this method tests an implementation detail of the method UpdateFileSystemOfRunner method
 	// if the implementation changes, delete this test and write a new one
 	copyRequest := &dto.UpdateFileSystemRequest{}
-	err := s.runner.UpdateFileSystem(copyRequest)
+	err := s.runner.UpdateFileSystem(copyRequest, context.Background())
 	s.NoError(err)
 	s.apiMock.AssertCalled(s.T(), "ExecuteCommand", mock.Anything, mock.Anything, mock.Anything,
 		false, mock.AnythingOfType("bool"), mock.Anything, mock.Anything, mock.Anything)
@@ -284,21 +284,21 @@ func (s *UpdateFileSystemTestSuite) TestUpdateFileSystemForRunnerPerformsTarExtr
 func (s *UpdateFileSystemTestSuite) TestUpdateFileSystemForRunnerReturnsErrorIfExitCodeIsNotZero() {
 	s.mockedExecuteCommandCall.Return(1, nil)
 	copyRequest := &dto.UpdateFileSystemRequest{}
-	err := s.runner.UpdateFileSystem(copyRequest)
+	err := s.runner.UpdateFileSystem(copyRequest, context.Background())
 	s.ErrorIs(err, ErrorFileCopyFailed)
 }
 
 func (s *UpdateFileSystemTestSuite) TestUpdateFileSystemForRunnerReturnsErrorIfApiCallDid() {
 	s.mockedExecuteCommandCall.Return(0, tests.ErrDefault)
 	copyRequest := &dto.UpdateFileSystemRequest{}
-	err := s.runner.UpdateFileSystem(copyRequest)
+	err := s.runner.UpdateFileSystem(copyRequest, context.Background())
 	s.ErrorIs(err, nomad.ErrorExecutorCommunicationFailed)
 }
 
 func (s *UpdateFileSystemTestSuite) TestFilesToCopyAreIncludedInTarArchive() {
 	copyRequest := &dto.UpdateFileSystemRequest{Copy: []dto.File{
 		{Path: tests.DefaultFileName, Content: []byte(tests.DefaultFileContent)}}}
-	err := s.runner.UpdateFileSystem(copyRequest)
+	err := s.runner.UpdateFileSystem(copyRequest, context.Background())
 	s.NoError(err)
 	s.apiMock.AssertCalled(s.T(), "ExecuteCommand", mock.Anything, mock.Anything, mock.Anything, false, true,
 		mock.Anything, mock.Anything, mock.Anything)
@@ -314,7 +314,7 @@ func (s *UpdateFileSystemTestSuite) TestFilesToCopyAreIncludedInTarArchive() {
 func (s *UpdateFileSystemTestSuite) TestTarFilesContainCorrectPathForRelativeFilePath() {
 	copyRequest := &dto.UpdateFileSystemRequest{Copy: []dto.File{
 		{Path: tests.DefaultFileName, Content: []byte(tests.DefaultFileContent)}}}
-	err := s.runner.UpdateFileSystem(copyRequest)
+	err := s.runner.UpdateFileSystem(copyRequest, context.Background())
 	s.Require().NoError(err)
 
 	tarFiles := s.readFilesFromTarArchive(s.stdin)
@@ -326,7 +326,7 @@ func (s *UpdateFileSystemTestSuite) TestTarFilesContainCorrectPathForRelativeFil
 func (s *UpdateFileSystemTestSuite) TestFilesWithAbsolutePathArePutInAbsoluteLocation() {
 	copyRequest := &dto.UpdateFileSystemRequest{Copy: []dto.File{
 		{Path: tests.FileNameWithAbsolutePath, Content: []byte(tests.DefaultFileContent)}}}
-	err := s.runner.UpdateFileSystem(copyRequest)
+	err := s.runner.UpdateFileSystem(copyRequest, context.Background())
 	s.Require().NoError(err)
 
 	tarFiles := s.readFilesFromTarArchive(s.stdin)
@@ -336,7 +336,7 @@ func (s *UpdateFileSystemTestSuite) TestFilesWithAbsolutePathArePutInAbsoluteLoc
 
 func (s *UpdateFileSystemTestSuite) TestDirectoriesAreMarkedAsDirectoryInTar() {
 	copyRequest := &dto.UpdateFileSystemRequest{Copy: []dto.File{{Path: tests.DefaultDirectoryName, Content: []byte{}}}}
-	err := s.runner.UpdateFileSystem(copyRequest)
+	err := s.runner.UpdateFileSystem(copyRequest, context.Background())
 	s.Require().NoError(err)
 
 	tarFiles := s.readFilesFromTarArchive(s.stdin)
@@ -349,7 +349,7 @@ func (s *UpdateFileSystemTestSuite) TestDirectoriesAreMarkedAsDirectoryInTar() {
 
 func (s *UpdateFileSystemTestSuite) TestFilesToRemoveGetRemoved() {
 	copyRequest := &dto.UpdateFileSystemRequest{Delete: []dto.FilePath{tests.DefaultFileName}}
-	err := s.runner.UpdateFileSystem(copyRequest)
+	err := s.runner.UpdateFileSystem(copyRequest, context.Background())
 	s.NoError(err)
 	s.apiMock.AssertCalled(s.T(), "ExecuteCommand", mock.Anything, mock.Anything, mock.Anything, false, true,
 		mock.Anything, mock.Anything, mock.Anything)
@@ -358,7 +358,7 @@ func (s *UpdateFileSystemTestSuite) TestFilesToRemoveGetRemoved() {
 
 func (s *UpdateFileSystemTestSuite) TestFilesToRemoveGetEscaped() {
 	copyRequest := &dto.UpdateFileSystemRequest{Delete: []dto.FilePath{"/some/potentially/harmful'filename"}}
-	err := s.runner.UpdateFileSystem(copyRequest)
+	err := s.runner.UpdateFileSystem(copyRequest, context.Background())
 	s.NoError(err)
 	s.apiMock.AssertCalled(s.T(), "ExecuteCommand", mock.Anything, mock.Anything, mock.Anything, false, true,
 		mock.Anything, mock.Anything, mock.Anything)
@@ -367,7 +367,7 @@ func (s *UpdateFileSystemTestSuite) TestFilesToRemoveGetEscaped() {
 
 func (s *UpdateFileSystemTestSuite) TestResetTimerGetsCalled() {
 	copyRequest := &dto.UpdateFileSystemRequest{}
-	err := s.runner.UpdateFileSystem(copyRequest)
+	err := s.runner.UpdateFileSystem(copyRequest, context.Background())
 	s.NoError(err)
 	s.timer.AssertCalled(s.T(), "ResetTimeout")
 }

internal/runner/runner.go

@@ -51,7 +51,7 @@ type Runner interface {
 
 	// UpdateFileSystem processes a dto.UpdateFileSystemRequest by first deleting each given dto.FilePath recursively
 	// and then copying each given dto.File to the runner.
-	UpdateFileSystem(request *dto.UpdateFileSystemRequest) error
+	UpdateFileSystem(request *dto.UpdateFileSystemRequest, ctx context.Context) error
 
 	// GetFileContent streams the file content at the requested path into the Writer provided at content.
 	// The result is streamed via the io.Writer in order to not overload the memory with user input.

internal/runner/runner_mock.go

@@ -1,4 +1,4 @@
-// Code generated by mockery v2.14.0. DO NOT EDIT.
+// Code generated by mockery v2.16.0. DO NOT EDIT.
 
 package runner
 
@@ -186,13 +186,13 @@ func (_m *RunnerMock) TimeoutPassed() bool {
 	return r0
 }
 
-// UpdateFileSystem provides a mock function with given fields: request
+// UpdateFileSystem provides a mock function with given fields: request, ctx
-func (_m *RunnerMock) UpdateFileSystem(request *dto.UpdateFileSystemRequest) error {
+func (_m *RunnerMock) UpdateFileSystem(request *dto.UpdateFileSystemRequest, ctx context.Context) error {
-	ret := _m.Called(request)
+	ret := _m.Called(request, ctx)
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(*dto.UpdateFileSystemRequest) error); ok {
+	if rf, ok := ret.Get(0).(func(*dto.UpdateFileSystemRequest, context.Context) error); ok {
-		r0 = rf(request)
+		r0 = rf(request, ctx)
 	} else {
 		r0 = ret.Error(0)
 	}