Implement merge request comments

environment/manager.go

@@ -1,7 +1,9 @@
 package environment
 
 import (
+	"fmt"
 	nomadApi "github.com/hashicorp/nomad/api"
+	"github.com/hashicorp/nomad/nomad/structs"
 	"gitlab.hpi.de/codeocean/codemoon/poseidon/api/dto"
 	"gitlab.hpi.de/codeocean/codemoon/poseidon/logging"
 	"gitlab.hpi.de/codeocean/codemoon/poseidon/nomad"
@@ -16,7 +18,7 @@ type Manager interface {
 	// CreateOrUpdate creates/updates an execution environment on the executor.
 	// Iff the job was created, the returned boolean is true and the returned error is nil.
 	CreateOrUpdate(
-		id string,
+		id runner.EnvironmentID,
 		request dto.ExecutionEnvironmentRequest,
 	) (bool, error)
 
@@ -24,9 +26,11 @@ type Manager interface {
 	Delete(id string)
 }
 
-func NewNomadEnvironmentManager(runnerManager runner.Manager, apiClient nomad.ExecutorAPI) *NomadEnvironmentManager {
+func NewNomadEnvironmentManager(runnerManager runner.Manager, apiClient nomad.ExecutorAPI) (
+	*NomadEnvironmentManager, error) {
 	environmentManager := &NomadEnvironmentManager{runnerManager, apiClient, *parseJob(defaultJobHCL)}
-	return environmentManager
+	err := environmentManager.loadExistingEnvironments()
+	return environmentManager, err
 }
 
 type NomadEnvironmentManager struct {
@@ -36,14 +40,10 @@ type NomadEnvironmentManager struct {
 }
 
 func (m *NomadEnvironmentManager) CreateOrUpdate(
-	id string,
+	id runner.EnvironmentID,
 	request dto.ExecutionEnvironmentRequest,
 ) (bool, error) {
-	idInt, err := strconv.Atoi(id)
-	if err != nil {
-		return false, err
-	}
-	err = m.registerDefaultJob(id,
+	templateJob, err := m.registerTemplateJob(id,
 		request.PrewarmingPoolSize, request.CPULimit, request.MemoryLimit,
 		request.Image, request.NetworkAccess, request.ExposedPorts)
 
@@ -51,7 +51,7 @@ func (m *NomadEnvironmentManager) CreateOrUpdate(
 		return false, err
 	}
 
-	created, err := m.runnerManager.CreateOrUpdateEnvironment(runner.EnvironmentID(idInt), request.PrewarmingPoolSize)
+	created, err := m.runnerManager.CreateOrUpdateEnvironment(id, request.PrewarmingPoolSize, templateJob)
 	if err != nil {
 		return created, err
 	}
@@ -61,3 +61,74 @@ func (m *NomadEnvironmentManager) CreateOrUpdate(
 func (m *NomadEnvironmentManager) Delete(id string) {
 
 }
+
+func (m *NomadEnvironmentManager) loadExistingEnvironments() error {
+	jobs, err := m.api.LoadAllJobs()
+	if err != nil {
+		return fmt.Errorf("can't load template jobs: %w", err)
+	}
+
+	var environmentTemplates, runnerJobs []*nomadApi.Job
+	for _, job := range jobs {
+		if nomad.IsEnvironmentTemplateID(*job.ID) {
+			environmentTemplates = append(environmentTemplates, job)
+		} else {
+			runnerJobs = append(runnerJobs, job)
+		}
+	}
+	m.recoverJobs(environmentTemplates, m.recoverEnvironmentTemplates)
+	m.recoverJobs(runnerJobs, m.recoverRunner)
+
+	err = m.runnerManager.ScaleAllEnvironments()
+	if err != nil {
+		return fmt.Errorf("can not restore environment scaling: %w", err)
+	}
+	return nil
+}
+
+type jobAdder func(id runner.EnvironmentID, job *nomadApi.Job, configTaskGroup *nomadApi.TaskGroup) error
+
+func (m *NomadEnvironmentManager) recoverEnvironmentTemplates(id runner.EnvironmentID, job *nomadApi.Job,
+	configTaskGroup *nomadApi.TaskGroup) error {
+	desiredIdleRunnersCount, err := strconv.Atoi(configTaskGroup.Meta[nomad.ConfigMetaPoolSizeKey])
+	if err != nil {
+		return fmt.Errorf("Couldn't convert pool size to int: %w", err)
+	}
+
+	m.runnerManager.RecoverEnvironment(id, job, uint(desiredIdleRunnersCount))
+	return nil
+}
+
+func (m *NomadEnvironmentManager) recoverRunner(id runner.EnvironmentID, job *nomadApi.Job,
+	configTaskGroup *nomadApi.TaskGroup) error {
+	isUsed := configTaskGroup.Meta[nomad.ConfigMetaUsedKey] == nomad.ConfigMetaUsedValue
+	m.runnerManager.RecoverRunner(id, job, isUsed)
+	return nil
+}
+
+func (m *NomadEnvironmentManager) recoverJobs(jobs []*nomadApi.Job, onJob jobAdder) {
+	for _, job := range jobs {
+		jobLogger := log.WithField("jobID", *job.ID)
+		if *job.Status != structs.JobStatusRunning {
+			jobLogger.Info("Job not running, skipping ...")
+			continue
+		}
+		configTaskGroup := nomad.FindConfigTaskGroup(job)
+		if configTaskGroup == nil {
+			jobLogger.Info("Couldn't find config task group in job, skipping ...")
+			continue
+		}
+		environmentID, err := runner.NewEnvironmentID(configTaskGroup.Meta[nomad.ConfigMetaEnvironmentKey])
+		if err != nil {
+			jobLogger.WithField("environmentID", configTaskGroup.Meta[nomad.ConfigMetaEnvironmentKey]).
+				WithError(err).
+				Error("Couldn't convert environment id of template job to int")
+			continue
+		}
+		err = onJob(environmentID, job, configTaskGroup)
+		if err != nil {
+			jobLogger.WithError(err).Info("Could not recover job.")
+			continue
+		}
+	}
+}