-
Notifications
You must be signed in to change notification settings - Fork 250
/
update-worker-pool.go
101 lines (87 loc) · 3.2 KB
/
update-worker-pool.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
package main
import (
"log"
"net/url"
"os"
"os/exec"
"path/filepath"
"strings"
"github.com/taskcluster/slugid-go/slugid"
tcclient "github.com/taskcluster/taskcluster/v47/clients/client-go"
"github.com/taskcluster/taskcluster/v47/clients/client-go/tcworkermanager"
)
func gitRevision(dir string) string {
command := exec.Command("git", "rev-parse", "HEAD")
command.Dir = dir
output, err := command.Output()
if err != nil {
panic(err)
}
return string(output[:len(output)-1])
}
func main() {
if len(os.Args) != 2 {
log.Fatal("Usage: " + os.Args[0] + " WORKER_TYPE_DIRECTORY")
}
workerTypeDir := os.Args[1]
absFile, err := filepath.Abs(workerTypeDir)
if err != nil {
log.Fatalf("File/directory '%v' could not be read due to '%s'", workerTypeDir, err)
}
workerType := filepath.Base(absFile)
workerPoolId := os.Getenv("PROVISIONER_ID") + "/" + workerType
files, err := os.ReadDir(workerTypeDir)
if err != nil {
log.Fatalf("File/directory '%v' (%v) could not be read due to '%s'", workerTypeDir, absFile, err)
}
scriptName := ""
for _, f := range files {
if !f.IsDir() && strings.HasPrefix(f.Name(), "bootstrap.") {
scriptName = f.Name()
}
}
workerManager := tcworkermanager.NewFromEnv()
cdv := tcclient.Client(*workerManager)
cd := &cdv
var wt map[string]interface{}
_, _, err = cd.APICall(nil, "GET", "/worker-pool/"+url.QueryEscape(workerPoolId), &wt, nil)
if err != nil {
log.Fatal(err)
}
workerPoolConfig := wt["config"].(map[string]interface{})
disks := workerPoolConfig["disks"].([]interface{})
delete(wt, "workerPoolId")
delete(wt, "created")
delete(wt, "lastModified")
diskObj := disks[0].(map[string]interface{})
initializeParams := diskObj["initializeParams"].(map[string]interface{})
oldImage := initializeParams["sourceImage"].(string)
newImage := "projects/pmoore-dev/global/images/" + os.Getenv("UNIQUE_NAME")
log.Print("Old image: " + oldImage)
log.Print("New image: " + newImage)
initializeParams["sourceImage"] = newImage
if oldImage == newImage {
log.Print("WARNING: No change to image (" + oldImage + ")")
} else {
log.Print("Worker pool " + workerPoolId + " updated to use image " + newImage)
}
userDataMap := workerPoolConfig["workerConfig"].(map[string]interface{})
genericWorker := userDataMap["genericWorker"].(map[string]interface{})
config := genericWorker["config"].(map[string]interface{})
oldDeploymentID := config["deploymentId"].(string)
newDeploymentID := slugid.Nice()
config["deploymentId"] = newDeploymentID
log.Print("Old deployment ID: " + oldDeploymentID)
log.Print("New deployment ID: " + newDeploymentID)
oldMetadata := config["workerTypeMetadata"].(map[string]interface{})
oldMachineSetup := oldMetadata["machine-setup"].(map[string]interface{})
oldScript := oldMachineSetup["script"].(string)
newScript := "https://raw.githubusercontent.com/taskcluster/generic-worker/" + gitRevision(workerTypeDir) + "/worker_types/" + workerType + "/" + scriptName
oldMachineSetup["script"] = newScript
log.Print("Old machine setup script: " + oldScript)
log.Print("New machine setup script: " + newScript)
_, _, err = cd.APICall(wt, "POST", "/worker-pool/"+url.QueryEscape(workerPoolId), new(interface{}), nil)
if err != nil {
log.Fatal(err)
}
}