feat: [CDE-344]: Skip CDE VMs in the dlite instance purger with the help of labels on the instance/VM

internal/drivers/distributed_manager.go

@@ -33,7 +33,7 @@ func (d *DistributedManager) BuildPools(ctx context.Context) error {
 // This helps in cleaning the pools
 func (d *DistributedManager) CleanPools(ctx context.Context, destroyBusy, destroyFree bool) error {
 	var returnError error
-	query := types.QueryParams{RunnerName: d.runnerName}
+	query := types.QueryParams{RunnerName: d.runnerName, MatchLabels: map[string]string{"retain": "false"}}
 	for _, pool := range d.poolMap {
 		err := d.cleanPool(ctx, pool, &query, destroyBusy, destroyFree)
 		if err != nil {
@@ -99,8 +99,9 @@ func (d *DistributedManager) StartInstancePurger(ctx context.Context, maxAgeBusy
 				case <-d.cleanupTimer.C:
 					logrus.Traceln("distributed dlite: Launching instance purger")
 
+					queryParams := types.QueryParams{MatchLabels: map[string]string{"retain": "false"}}
 					for _, pool := range d.poolMap {
-						if err := d.startInstancePurger(ctx, pool, maxAgeBusy); err != nil {
+						if err := d.startInstancePurger(ctx, pool, maxAgeBusy, queryParams); err != nil {
 							logger.FromContext(ctx).WithError(err).
 								Errorln("distributed dlite: purger: Failed to purge stale instances")
 						}
@@ -113,7 +114,7 @@ func (d *DistributedManager) StartInstancePurger(ctx context.Context, maxAgeBusy
 	return nil
 }
 
-func (d *DistributedManager) startInstancePurger(ctx context.Context, pool *poolEntry, maxAgeBusy time.Duration) error {
+func (d *DistributedManager) startInstancePurger(ctx context.Context, pool *poolEntry, maxAgeBusy time.Duration, queryParams types.QueryParams) error {
 	logr := logger.FromContext(ctx).
 		WithField("driver", pool.Driver.DriverName()).
 		WithField("pool", pool.Name)
@@ -130,6 +131,10 @@ func (d *DistributedManager) startInstancePurger(ctx context.Context, pool *pool
 			squirrel.Eq{"instance_state": types.StateInUse},
 			squirrel.Lt{"instance_started": currentTime.Add(-maxAgeBusy).Unix()},
 		}
+		for key, value := range queryParams.MatchLabels {
+			condition := squirrel.Expr("(instance_labels->>?) = ?", key, value)
+			busyCondition = append(busyCondition, condition)
+		}
 		conditions = append(conditions, busyCondition)
 	}
 

internal/drivers/manager.go

@@ -252,7 +252,8 @@ func (m *Manager) StartInstancePurger(ctx context.Context, maxAgeBusy, maxAgeFre
 							pool.Lock()
 							defer pool.Unlock()
 
-							busy, free, hibernating, err := m.List(ctx, pool, nil)
+							queryParams := &types.QueryParams{MatchLabels: map[string]string{"retain": "false"}}
+							busy, free, hibernating, err := m.List(ctx, pool, queryParams)
 							if err != nil {
 								return fmt.Errorf("failed to list instances of pool=%q error: %w", pool.Name, err)
 							}
@@ -565,8 +566,18 @@ func (m *Manager) buildPoolWithMutex(ctx context.Context, pool *poolEntry, tlsSe
 	return m.buildPool(ctx, pool, tlsServerName, query)
 }
 
-func (m *Manager) setupInstance(ctx context.Context, pool *poolEntry, tlsServerName, ownerID, resourceClass string, inuse bool, agentConfig *types.GitspaceAgentConfig, storageConfig *types.StorageConfig) (*types.Instance, error) {
+func (m *Manager) setupInstance(
+	ctx context.Context,
+	pool *poolEntry,
+	tlsServerName,
+	ownerID,
+	resourceClass string,
+	inuse bool,
+	agentConfig *types.GitspaceAgentConfig,
+	storageConfig *types.StorageConfig,
+) (*types.Instance, error) {
 	var inst *types.Instance
+	retain := "false"
 
 	// generate certs
 	createOptions, err := certs.Generate(m.runnerName, tlsServerName)
@@ -589,13 +600,15 @@ func (m *Manager) setupInstance(ctx context.Context, pool *poolEntry, tlsServerN
 		}
 	}
 	createOptions.AutoInjectionBinaryURI = m.autoInjectionBinaryURI
-	if agentConfig != nil {
+	if agentConfig != nil && agentConfig.Secret != "" {
 		createOptions.GitspaceOpts = types.GitspaceOpts{
 			Secret:      agentConfig.Secret,
 			AccessToken: agentConfig.AccessToken,
 			Ports:       agentConfig.Ports,
 		}
+		retain = "true"
 	}
+	createOptions.Labels = map[string]string{"retain": retain}
 	if err != nil {
 		logrus.WithError(err).
 			Errorln("manager: failed to generate certificates")

internal/drivers/nomad/driver.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	_ "embed"
 	"encoding/base64"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"net"
@@ -234,6 +235,10 @@ func (p *config) Create(ctx context.Context, opts *types.InstanceCreateOpts) (*t
 		logr = logr.WithField("gitspaces_port_mapping", gitspacesPortMappingsString)
 	}
 
+	labelsBytes, marshalErr := json.Marshal(opts.Labels)
+	if marshalErr != nil {
+		return nil, fmt.Errorf("scheduler: could not marshal labels: %v, err: %w", opts.Labels, marshalErr)
+	}
 	instance := &types.Instance{
 		ID:                   vm,
 		NodeID:               id,
@@ -251,7 +256,10 @@ func (p *config) Create(ctx context.Context, opts *types.InstanceCreateOpts) (*t
 		Port:                 int64(liteEngineHostPort),
 		GitspacePortMappings: gitspacesPortMappings,
 		Address:              ip,
-		StorageIdentifier:    opts.StorageOpts.CephPoolIdentifier + "/" + opts.StorageOpts.Identifier,
+		Labels:               labelsBytes,
+	}
+	if opts.StorageOpts.Identifier != "" {
+		instance.StorageIdentifier = opts.StorageOpts.CephPoolIdentifier + "/" + opts.StorageOpts.Identifier
 	}
 
 	logr.Debugln("scheduler: submitting VM creation job")
@@ -486,13 +494,14 @@ func (p *config) destroyJob(ctx context.Context, vm, nodeID, storageIdentifier s
 		},
 	}
 	if storageIdentifier != "" {
-		job.TaskGroups[0].Tasks = append(job.TaskGroups[0].Tasks, p.getCephStorageScriptCreateTask(cephStorageScriptEncoded, cephStorageScriptPath))
-		job.TaskGroups[0].Tasks = append(job.TaskGroups[0].Tasks, p.getCephStorageScriptCleanupTask(cephStorageScriptPath))
+		job.TaskGroups[0].Tasks = append(job.TaskGroups[0].Tasks,
+			p.getCephStorageScriptCreateTask(cephStorageScriptEncoded, cephStorageScriptPath),
+			p.getCephStorageScriptCleanupTask(cephStorageScriptPath))
 	}
 	return job, id
 }
 
-func cleanupStorage(vm string, storageIdentifier string, storageCleanupType *storage.CleanupType, destroyCmd *string) (string, string, error) {
+func cleanupStorage(vm, storageIdentifier string, storageCleanupType *storage.CleanupType, destroyCmd *string) (cephStorageScriptEncoded, cephStorageScriptPath string, err error) {
 	var cephStorageCleanupScriptTemplate *template.Template
 	if *storageCleanupType == storage.Detach {
 		cephStorageCleanupScriptTemplate = template.Must(template.New("detach-ceph-storage").Funcs(funcs).Parse(detachCephStorageScript))
@@ -502,6 +511,7 @@ func cleanupStorage(vm string, storageIdentifier string, storageCleanupType *sto
 
 	sb := &strings.Builder{}
 	storageIdentifierSplit := strings.Split(storageIdentifier, "/")
+	//nolint:gomnd
 	if len(storageIdentifierSplit) != 2 {
 		return "", "", fmt.Errorf("scheduler: could not parse storage identifier %s", storageIdentifier)
 	}
@@ -512,12 +522,12 @@ func cleanupStorage(vm string, storageIdentifier string, storageCleanupType *sto
 		CephPoolIdentifier: storageIdentifierSplit[0],
 		RBDIdentifier:      storageIdentifierSplit[1],
 	}
-	err := cephStorageCleanupScriptTemplate.Execute(sb, params)
+	err = cephStorageCleanupScriptTemplate.Execute(sb, params)
 	if err != nil {
 		return "", "", fmt.Errorf("scheduler: failed to execute de-provision-ceph-storage template to get the script: %w", err)
 	}
-	cephStorageScriptEncoded := base64.StdEncoding.EncodeToString([]byte(sb.String()))
-	cephStorageScriptPath := fmt.Sprintf("/usr/local/bin/%s_delete_ceph_storage.sh", vm)
+	cephStorageScriptEncoded = base64.StdEncoding.EncodeToString([]byte(sb.String()))
+	cephStorageScriptPath = fmt.Sprintf("/usr/local/bin/%s_delete_ceph_storage.sh", vm)
 	*destroyCmd += fmt.Sprintf("\ncat %s | base64 --decode | bash", cephStorageScriptPath)
 	return cephStorageScriptEncoded, cephStorageScriptPath, nil
 }
@@ -673,7 +683,7 @@ func (p *config) getCephStorageScriptCleanupTask(deProvisionCephStorageScriptPat
 	}
 }
 
-func (p *config) getCephStorageScriptCreateTask(cephStorageScriptEncoded string, cephStorageScriptPath string) *api.Task {
+func (p *config) getCephStorageScriptCreateTask(cephStorageScriptEncoded, cephStorageScriptPath string) *api.Task {
 	return &api.Task{
 		Name:      "create_ceph_storage_cleanup_script_on_host",
 		Driver:    "raw_exec",

internal/drivers/nomad/linux_virtualizer.go

@@ -45,6 +45,14 @@ func (lv *LinuxVirtualizer) GetInitJob(vm, nodeID, vmImage, userData, username,
 	runCmdFormat = "%s run %s --name %s --cpus %s --memory %sGB --size %s --ssh --runtime=docker --ports %d:%s --copy-files %s:%s"
 	args := []interface{}{ignitePath, vmImage, vm, resource.Cpus, resource.MemoryGB, resource.DiskSize, port, strconv.Itoa(lehelper.LiteEnginePort), hostPath, vmPath}
 
+	// add labels
+	if len(opts.Labels) > 0 {
+		runCmdFormat += " --label"
+		for key, value := range opts.Labels {
+			runCmdFormat += fmt.Sprintf(" %s=%s", key, value)
+		}
+	}
+
 	// gitspace args
 	for vmPort, hostPort := range gitspacesPortMappings {
 		runCmdFormat += " --ports %d:%d"
@@ -279,7 +287,7 @@ func (lv *LinuxVirtualizer) GetDestroyScriptGenerator() func(string) string {
 	}
 }
 
-func (lv *LinuxVirtualizer) getScriptCleanupCmd(opts *types.InstanceCreateOpts, hostPath string, provisionCephStorageScriptPath string) string {
+func (lv *LinuxVirtualizer) getScriptCleanupCmd(opts *types.InstanceCreateOpts, hostPath, provisionCephStorageScriptPath string) string {
 	cleanUpCmdFormat := "rm %s"
 	cleanUpCmdArgs := []interface{}{hostPath}
 	if opts.StorageOpts.Identifier != "" && provisionCephStorageScriptPath != "" {

store/database/ldb/instances.go

@@ -4,11 +4,13 @@ import (
 	"bytes"
 	"context"
 	"encoding/gob"
+	"encoding/json"
 	"sort"
 	"time"
 
 	"github.com/drone-runners/drone-runner-aws/store"
 	"github.com/drone-runners/drone-runner-aws/types"
+	"github.com/sirupsen/logrus"
 	"github.com/syndtr/goleveldb/leveldb"
 	"github.com/syndtr/goleveldb/leveldb/util"
 )
@@ -102,6 +104,7 @@ func (s InstanceStore) DeleteAndReturn(ctx context.Context, query string, args .
 }
 
 func (s InstanceStore) satisfy(inst *types.Instance, pool string, params *types.QueryParams) bool {
+	log := logrus.New()
 	if pool == "" {
 		return true
 	}
@@ -121,6 +124,19 @@ func (s InstanceStore) satisfy(inst *types.Instance, pool string, params *types.
 				return false
 			}
 		}
+		if len(params.MatchLabels) > 0 {
+			var instanceLabels map[string]string
+			err := json.Unmarshal(inst.Labels, &instanceLabels)
+			if err != nil {
+				log.Errorln("Error decoding instance labels json:", err)
+				return false
+			}
+			for key, value := range params.MatchLabels {
+				if instVal, ok := instanceLabels[key]; !ok || instVal != value {
+					return false
+				}
+			}
+		}
 	}
 	return true
 }

store/database/migrate/postgres/0010_update_table_instances_labels.up.sql

@@ -0,0 +1 @@
+ALTER TABLE instances ADD COLUMN instance_labels JSONB NOT NULL DEFAULT '{"retain":"false"}';

store/database/migrate/sqlite/0008_update_table_instances_labels.up.sql

@@ -0,0 +1 @@
+ALTER TABLE instances ADD COLUMN instance_labels JSONB NOT NULL DEFAULT '{"retain":"false"}';

store/database/sql/instances.go

@@ -51,6 +51,11 @@ func (s InstanceStore) List(_ context.Context, pool string, params *types.QueryP
 			stmt = stmt.Where(squirrel.Eq{"runner_name": params.RunnerName})
 			args = append(args, params.RunnerName)
 		}
+		for key, value := range params.MatchLabels {
+			condition := squirrel.Expr("(instance_labels->>?) = ?", key, value)
+			stmt = stmt.Where(condition)
+			args = append(args, key, value)
+		}
 	}
 	stmt = stmt.OrderBy("instance_started " + "ASC")
 	sql, _, _ := stmt.ToSql()
@@ -149,6 +154,7 @@ const instanceColumns = `
 ,instance_port
 ,instance_owner_id
 ,instance_storage_identifier
+,instance_labels
 `
 
 const instanceFindByID = `SELECT ` + instanceColumns + `
@@ -186,6 +192,7 @@ INSERT INTO instances (
 ,instance_owner_id
 ,runner_name
 ,instance_storage_identifier
+,instance_labels
 ) values (
  :instance_id
 ,:instance_node_id
@@ -215,6 +222,7 @@ INSERT INTO instances (
 ,:instance_owner_id
 ,:runner_name
 ,:instance_storage_identifier
+,:instance_labels
 ) RETURNING instance_id
 `
 

types/types.go

@@ -67,6 +67,7 @@ type Instance struct {
 	RunnerName           string      `db:"runner_name" json:"runner_name"`
 	GitspacePortMappings map[int]int `json:"gitspaces_port_mappings"`
 	StorageIdentifier    string      `db:"instance_storage_identifier" json:"storage_identifier"`
+	Labels               []byte      `db:"instance_labels" json:"instance_labels"`
 }
 
 type Tmate struct {
@@ -98,6 +99,7 @@ type InstanceCreateOpts struct {
 	GitspaceOpts           GitspaceOpts
 	StorageOpts            StorageOpts
 	AutoInjectionBinaryURI string
+	Labels                 map[string]string
 }
 
 // Platform defines the target platform.
@@ -110,10 +112,11 @@ type Platform struct {
 }
 
 type QueryParams struct {
-	Status     InstanceState
-	Stage      string
-	Platform   *Platform
-	RunnerName string
+	Status      InstanceState
+	Stage       string
+	Platform    *Platform
+	RunnerName  string
+	MatchLabels map[string]string
 }
 
 type StageOwner struct {