diff --git a/internal/runner/abstract_manager.go b/internal/runner/abstract_manager.go index 6fa06fe..8e9c9cd 100644 --- a/internal/runner/abstract_manager.go +++ b/internal/runner/abstract_manager.go @@ -64,10 +64,14 @@ func (n *AbstractManager) EnvironmentStatistics() map[dto.EnvironmentID]*dto.Sta ID: int(e.ID()), PrewarmingPoolSize: e.PrewarmingPoolSize(), IdleRunners: uint(e.IdleRunnerCount()), - UsedRunners: 0, // Adjusted later as the information is present in the specific runner manager. + UsedRunners: 0, // Increased below. } } + for _, r := range n.usedRunners.List() { + environments[r.Environment()].UsedRunners++ + } + return environments } diff --git a/internal/runner/aws_manager.go b/internal/runner/aws_manager.go index 5487779..d7889a3 100644 --- a/internal/runner/aws_manager.go +++ b/internal/runner/aws_manager.go @@ -45,22 +45,3 @@ func (a AWSRunnerManager) Return(r Runner) error { } return nil } - -// EnvironmentStatistics returns only the used runner for each environment as the prewarming is handled -// by AWS transparently. -func (a AWSRunnerManager) EnvironmentStatistics() map[dto.EnvironmentID]*dto.StatisticalExecutionEnvironmentData { - environments := a.AbstractManager.EnvironmentStatistics() - - for _, r := range a.usedRunners.List() { - workload, isAWSRunner := r.(*AWSFunctionWorkload) - if !isAWSRunner { - log.WithField("workload", workload).Error("Stored runners must be AWS runner") - continue - } - - environmentID := workload.environment.ID() - environments[environmentID].UsedRunners++ - } - - return environments -} diff --git a/internal/runner/aws_runner.go b/internal/runner/aws_runner.go index 6f2641d..318c456 100644 --- a/internal/runner/aws_runner.go +++ b/internal/runner/aws_runner.go @@ -61,6 +61,10 @@ func (w *AWSFunctionWorkload) ID() string { return w.id } +func (w *AWSFunctionWorkload) Environment() dto.EnvironmentID { + return w.environment.ID() +} + func (w *AWSFunctionWorkload) MappedPorts() []*dto.MappedPort { return []*dto.MappedPort{} } diff --git a/internal/runner/nomad_manager.go b/internal/runner/nomad_manager.go index e971d41..c9f77ba 100644 --- a/internal/runner/nomad_manager.go +++ b/internal/runner/nomad_manager.go @@ -34,19 +34,6 @@ func NewNomadRunnerManager(apiClient nomad.ExecutorAPI, ctx context.Context) *No return m } -func (m *NomadRunnerManager) EnvironmentStatistics() map[dto.EnvironmentID]*dto.StatisticalExecutionEnvironmentData { - environments := m.AbstractManager.EnvironmentStatistics() - - for _, r := range m.usedRunners.List() { - id, err := nomad.EnvironmentIDFromRunnerID(r.ID()) - if err != nil { - log.WithError(err).Error("Stored runners must have correct IDs") - } - environments[id].UsedRunners++ - } - return environments -} - func (m *NomadRunnerManager) Claim(environmentID dto.EnvironmentID, duration int) (Runner, error) { environment, ok := m.environments.Get(environmentID) if !ok { diff --git a/internal/runner/nomad_runner.go b/internal/runner/nomad_runner.go index 3a79eef..bb01090 100644 --- a/internal/runner/nomad_runner.go +++ b/internal/runner/nomad_runner.go @@ -66,6 +66,14 @@ func (r *NomadJob) ID() string { return r.id } +func (r *NomadJob) Environment() dto.EnvironmentID { + id, err := nomad.EnvironmentIDFromRunnerID(r.ID()) + if err != nil { + log.WithError(err).Error("Runners must have correct IDs") + } + return id +} + func (r *NomadJob) MappedPorts() []*dto.MappedPort { ports := make([]*dto.MappedPort, 0, len(r.portMappings)) for _, portMapping := range r.portMappings { diff --git a/internal/runner/runner.go b/internal/runner/runner.go index cd09de8..066b784 100644 --- a/internal/runner/runner.go +++ b/internal/runner/runner.go @@ -19,6 +19,9 @@ type Runner interface { // ID returns the id of the runner. ID() string + // Environment returns the id of the Environment to which the Runner belongs. + Environment() dto.EnvironmentID + // MappedPorts returns the mapped ports of the runner. MappedPorts() []*dto.MappedPort diff --git a/internal/runner/runner_mock.go b/internal/runner/runner_mock.go index f44aa0f..5879ceb 100644 --- a/internal/runner/runner_mock.go +++ b/internal/runner/runner_mock.go @@ -1,4 +1,4 @@ -// Code generated by mockery v0.0.0-dev. DO NOT EDIT. +// Code generated by mockery v2.10.4. DO NOT EDIT. package runner @@ -32,6 +32,20 @@ func (_m *RunnerMock) Destroy() error { return r0 } +// Environment provides a mock function with given fields: +func (_m *RunnerMock) Environment() dto.EnvironmentID { + ret := _m.Called() + + var r0 dto.EnvironmentID + if rf, ok := ret.Get(0).(func() dto.EnvironmentID); ok { + r0 = rf() + } else { + r0 = ret.Get(0).(dto.EnvironmentID) + } + + return r0 +} + // ExecuteInteractively provides a mock function with given fields: id, stdin, stdout, stderr func (_m *RunnerMock) ExecuteInteractively(id string, stdin io.ReadWriter, stdout io.Writer, stderr io.Writer) (<-chan ExitInfo, context.CancelFunc, error) { ret := _m.Called(id, stdin, stdout, stderr)