diff --git a/src/lib/config/config.go b/src/lib/config/config.go index b9158fd..215084d 100644 --- a/src/lib/config/config.go +++ b/src/lib/config/config.go @@ -152,14 +152,21 @@ type GitHubOrganization struct { } type TrayType struct { - Name string `yaml:"name" validate:"required"` - Provider string `yaml:"provider" validate:"required"` - RunnerGroupId int64 `yaml:"runnerGroupId" validate:"required"` - Shutdown bool `yaml:"shutdown"` - GitHubOrg string `yaml:"githubOrg" validate:"required"` - MaxTrays int `yaml:"limit"` - Config TrayConfig `yaml:"config"` - ExtraMetadata TrayExtraMetadata + Name string `yaml:"name" validate:"required"` + Provider string `yaml:"provider" validate:"required"` + RunnerGroupId int64 `yaml:"runnerGroupId" validate:"required"` + Shutdown bool `yaml:"shutdown"` + GitHubOrg string `yaml:"githubOrg" validate:"required"` + MaxTrays int `yaml:"limit"` + Config TrayConfig `yaml:"config"` + ExtraMetadata TrayExtraMetadata + DeleteMetadata []string `yaml:"deleteMetadata"` +} + +func (t *TrayType) GetMetadataKeysToDelete() []string { + keys := make([]string, 0) + keys = append(keys, t.DeleteMetadata...) + return keys } type TrayExtraMetadata map[string]string diff --git a/src/lib/trays/providers/dockerProvider.go b/src/lib/trays/providers/dockerProvider.go index 102b8a3..8a3bc2e 100644 --- a/src/lib/trays/providers/dockerProvider.go +++ b/src/lib/trays/providers/dockerProvider.go @@ -71,6 +71,10 @@ func (d *DockerProvider) RunTray(tray *trays.Tray) error { return nil } +func (d *DockerProvider) DeleteMetadata(tray *trays.Tray, keys []string) error { + return nil +} + func (d *DockerProvider) CleanTray(tray *trays.Tray) error { var dockerCommand = exec.Command("docker", "container", "stop", tray.GetId()) dockerCommandOutput, err := dockerCommand.CombinedOutput() diff --git a/src/lib/trays/providers/gceProvider.go b/src/lib/trays/providers/gceProvider.go index 8076371..c078706 100644 --- a/src/lib/trays/providers/gceProvider.go +++ b/src/lib/trays/providers/gceProvider.go @@ -157,6 +157,92 @@ func (g *GceProvider) createInstancesClient() (*compute.InstancesClient, error) return instancesClient, err } +func (g *GceProvider) DeleteMetadata(tray *trays.Tray, keys []string) error { + if len(keys) == 0 { + return nil + } + + ctx := context.Background() + + client, err := g.createInstancesClient() + if err != nil { + return err + } + + var ( + zone = tray.ProviderData["zone"] + project = g.providerConfig.Get("project") + ) + + // Get current metadata + instance, err := client.Get(ctx, &computepb.GetInstanceRequest{ + Project: project, + Zone: zone, + Instance: tray.GetId(), + }) + if err != nil { + return fmt.Errorf("failed to get instance metadata: %w", err) + } + + currentMetadata := instance.GetMetadata() + if currentMetadata == nil { + g.logger.Infof("Instance %s has no metadata, nothing to delete", tray.GetId()) + return nil + } + + // Build set of keys to delete + keysToDelete := make(map[string]bool, len(keys)) + for _, k := range keys { + keysToDelete[k] = true + } + + // Filter out the keys to delete + newItems, deletedKeys := filterMetadataKeys(currentMetadata.GetItems(), keysToDelete) + + if len(deletedKeys) == 0 { + g.logger.Infof("None of the specified metadata keys were found on instance %s", tray.GetId()) + return nil + } + + // Set the new metadata with the same fingerprint + op, err := client.SetMetadata(ctx, &computepb.SetMetadataInstanceRequest{ + Project: project, + Zone: zone, + Instance: tray.GetId(), + MetadataResource: &computepb.Metadata{ + Fingerprint: currentMetadata.Fingerprint, + Items: newItems, + }, + }) + if err != nil { + return fmt.Errorf("failed to set metadata: %w", err) + } + + err = op.Wait(ctx) + if err != nil { + return fmt.Errorf("metadata update operation failed: %w", err) + } + + g.logger.Infof("Deleted metadata keys %v from instance %s", deletedKeys, tray.GetId()) + return nil +} + +func filterMetadataKeys(items []*computepb.Items, keysToDelete map[string]bool) ([]*computepb.Items, []string) { + var newItems []*computepb.Items + var deletedKeys []string + + for _, item := range items { + key := item.GetKey() + if keysToDelete[key] { + deletedKeys = append(deletedKeys, key) + continue + } + newItems = append(newItems, item) // build new list here + } + + return newItems, deletedKeys +} + func createGcpMetadata(fieldMaps ...map[string]string) *computepb.Metadata { var items []*computepb.Items diff --git a/src/lib/trays/providers/iTrayProvider.go b/src/lib/trays/providers/iTrayProvider.go index f4b3d98..2a43d72 100644 --- a/src/lib/trays/providers/iTrayProvider.go +++ b/src/lib/trays/providers/iTrayProvider.go @@ -18,4 +18,7 @@ type ITrayProvider interface { // CleanTray deletes the tray with the given ID. CleanTray(tray *trays.Tray) error + + // DeleteMetadata removes the specified metadata keys from the tray's instance. + DeleteMetadata(tray *trays.Tray, keys []string) error } diff --git a/src/server/handlers/agentHandler.go b/src/server/handlers/agentHandler.go index 589b5cf..81f684e 100644 --- a/src/server/handlers/agentHandler.go +++ b/src/server/handlers/agentHandler.go @@ -7,6 +7,7 @@ import ( "cattery/lib/messages" "cattery/lib/metrics" "cattery/lib/trays" + "cattery/lib/trays/providers" "encoding/json" "fmt" "net/http" @@ -110,6 +111,21 @@ func AgentRegister(responseWriter http.ResponseWriter, r *http.Request) { metrics.RegisteredTraysAdd(tray.GitHubOrgName, tray.TrayTypeName, 1) logger.Infof("Agent %s registered with runner ID %d", agentId, newAgent.RunnerId) + + // Delete instance metadata keys after registration + keysToDelete := trayType.GetMetadataKeysToDelete() + if len(keysToDelete) > 0 { + go func() { + provider, err := providers.GetProviderForTray(tray) + if err != nil { + logger.Warnf("Failed to get provider for metadata deletion: %v", err) + return + } + if err := provider.DeleteMetadata(tray, keysToDelete); err != nil { + logger.Warnf("Failed to delete instance metadata: %v", err) + } + }() + } } // validateAgentId validates the agent ID