func ensureEnvironment(t *testing.T) (*runtime.Environment, engines.Engine, plugins.Plugin) {
tempPath := filepath.Join(os.TempDir(), slugid.Nice())
tempStorage, err := runtime.NewTemporaryStorage(tempPath)
if err != nil {
t.Fatal(err)
}
environment := &runtime.Environment{
TemporaryStorage: tempStorage,
}
engineProvider := engines.Engines()["mock"]
engine, err := engineProvider.NewEngine(engines.EngineOptions{
Environment: environment,
Log: logger.WithField("engine", "mock"),
})
if err != nil {
t.Fatal(err.Error())
}
pluginOptions := plugins.PluginOptions{
Environment: environment,
Engine: engine,
Log: logger.WithField("component", "Plugin Manager"),
}
pm, err := plugins.Plugins()["success"].NewPlugin(pluginOptions)
if err != nil {
t.Fatalf("Error creating task manager. Could not create plugin manager. %s", err)
}
return environment, engine, pm
}
func (p *EngineProvider) ensureEngine() {
p.m.Lock()
defer p.m.Unlock()
p.refCount++
if p.engine != nil {
return
}
// Create a runtime environment
p.environment = newTestEnvironment()
// Find EngineProvider
engineProvider := engines.Engines()[p.Engine]
if engineProvider == nil {
fmtPanic("Couldn't find EngineProvider: ", p.Engine)
}
var jsonConfig interface{}
err := json.Unmarshal([]byte(p.Config), &jsonConfig)
nilOrPanic(err, "Config parsing failed: ", p.Config)
err = engineProvider.ConfigSchema().Validate(jsonConfig)
nilOrPanic(err, "Config validation failed: ", p.Config, "\nError: ", err)
// Create Engine instance
engine, err := engineProvider.NewEngine(engines.EngineOptions{
Environment: p.environment,
Log: p.environment.Log.WithField("engine", p.Engine),
Config: jsonConfig,
})
nilOrPanic(err, "Failed to create Engine")
p.engine = engine
}
// New will create a worker given configuration matching the schema from
// ConfigSchema(). The log parameter is optional and if nil is given a default
// logrus logger will be used.
func New(config interface{}, log *logrus.Logger) (*Worker, error) {
// Validate and map configuration to c
var c configType
if err := schematypes.MustMap(ConfigSchema(), config, &c); err != nil {
return nil, fmt.Errorf("Invalid configuration: %s", err)
}
// Create temporary folder
err := os.RemoveAll(c.TemporaryFolder)
if err != nil {
return nil, fmt.Errorf("Failed to remove temporaryFolder: %s, error: %s",
c.TemporaryFolder, err)
}
tempStorage, err := runtime.NewTemporaryStorage(c.TemporaryFolder)
if err != nil {
return nil, fmt.Errorf("Failed to create temporary folder, error: %s", err)
}
// Create logger
if log == nil {
log = logrus.New()
}
log.Level, _ = logrus.ParseLevel(c.LogLevel)
// Setup WebHookServer
localServer, err := webhookserver.NewLocalServer(
net.ParseIP(c.ServerIP), c.ServerPort,
c.NetworkInterface, c.ExposedPort,
c.DNSDomain,
c.DNSSecret,
c.TLSCertificate,
c.TLSKey,
time.Duration(c.MaxLifeCycle)*time.Second,
)
if err != nil {
return nil, err
}
// Create environment
gc := gc.New(c.TemporaryFolder, c.MinimumDiskSpace, c.MinimumMemory)
env := &runtime.Environment{
GarbageCollector: gc,
Log: log,
TemporaryStorage: tempStorage,
WebHookServer: localServer,
}
// Ensure that engine confiuguration was provided for the engine selected
if _, ok := c.Engines[c.Engine]; !ok {
return nil, fmt.Errorf("Invalid configuration: The key 'engines.%s' must "+
"be specified when engine '%s' is selected", c.Engine, c.Engine)
}
// Find engine provider (schema should ensure it exists)
provider := engines.Engines()[c.Engine]
engine, err := provider.NewEngine(engines.EngineOptions{
Environment: env,
Log: env.Log.WithField("engine", c.Engine),
Config: c.Engines[c.Engine],
})
if err != nil {
return nil, fmt.Errorf("Engine initialization failed, error: %s", err)
}
// Initialize plugin manager
pm, err := plugins.NewPluginManager(plugins.PluginOptions{
Environment: env,
Engine: engine,
Log: env.Log.WithField("plugin", "plugin-manager"),
Config: c.Plugins,
})
if err != nil {
return nil, fmt.Errorf("Plugin initialization failed, error: %s", err)
}
tm, err := newTaskManager(
&c, engine, pm, env,
env.Log.WithField("component", "task-manager"), gc,
)
if err != nil {
return nil, err
}
return &Worker{
log: env.Log.WithField("component", "worker"),
tm: tm,
sm: runtime.NewShutdownManager("local"),
env: env,
server: localServer,
done: make(chan struct{}),
}, nil
}
// ConfigSchema returns the configuration schema for the worker.
func ConfigSchema() schematypes.Object {
engineConfig := schematypes.Properties{}
engineNames := []string{}
for name, provider := range engines.Engines() {
engineNames = append(engineNames, name)
engineConfig[name] = provider.ConfigSchema()
}
return schematypes.Object{
MetaData: schematypes.MetaData{
Title: "Worker Configuration",
Description: `This contains configuration for the worker process.`,
},
Properties: schematypes.Properties{
"engine": schematypes.StringEnum{
MetaData: schematypes.MetaData{
Title: "Worker Engine",
Description: `Selected worker engine to use, notice that the
configuration for this engine **must** be present under the
'engines.<engine>' configuration key.`,
},
Options: engineNames,
},
"engines": schematypes.Object{
MetaData: schematypes.MetaData{
Title: "Engine Configuration",
Description: `Mapping from engine name to engine configuration.
Even-though the worker will only use one engine at any given time,
the configuration file can hold configuration for all engines.
Hence, you need only update the 'engine' key to change which engine
should be used.`,
},
Properties: engineConfig,
},
"plugins": plugins.PluginManagerConfigSchema(),
"capacity": schematypes.Integer{
MetaData: schematypes.MetaData{
Title: "Capacity",
Description: `The number of tasks that this worker supports running in
parallel.`,
},
Minimum: 1,
Maximum: 1000,
},
"credentials": schematypes.Object{
MetaData: schematypes.MetaData{
Title: "TaskCluster Credentials",
Description: `The set of credentials that should be used by the worker
when authenticating against taskcluster endpoints. This needs scopes
for claiming tasks for the given workerType.`,
},
Properties: schematypes.Properties{
"clientId": schematypes.String{
MetaData: schematypes.MetaData{
Title: "ClientId",
Description: `ClientId for credentials`,
},
Pattern: `^[A-Za-z0-9@/:._-]+$`,
},
"accessToken": schematypes.String{
MetaData: schematypes.MetaData{
Title: "AccessToken",
Description: `The security-sensitive access token for the client.`,
},
Pattern: `^[a-zA-Z0-9_-]{22,66}$`,
},
"certificate": schematypes.String{
MetaData: schematypes.MetaData{
Title: "Certificate",
Description: `The certificate for the client, if using temporary
credentials.`,
},
},
},
Required: []string{"clientId", "accessToken"},
},
"pollingInterval": schematypes.Integer{
MetaData: schematypes.MetaData{
Title: "Task Polling Interval",
Description: `The amount of time to wait between task polling
iterations in seconds.`,
},
Minimum: 0,
Maximum: 10 * 60,
},
"reclaimOffset": schematypes.Integer{
MetaData: schematypes.MetaData{
Title: "Reclaim Offset",
Description: `The number of seconds priorty task claim expiration the
claim should be reclamed.`,
},
Minimum: 0,
Maximum: 10 * 60,
},
"queueBaseUrl": schematypes.URI{
MetaData: schematypes.MetaData{
Title: "Queue BaseUrl",
Description: `BaseUrl for taskcluster-queue, defaults to value from the
taskcluster client library.`,
},
},
"provisionerId": schematypes.String{
MetaData: schematypes.MetaData{
Title: "ProvisionerId",
Description: `ProvisionerId for workerType that tasks should be claimed
from. Note, a 'workerType' is only unique given the 'provisionerId'.`,
},
Pattern: `^[a-zA-Z0-9_-]{1,22}$`,
},
"workerType": schematypes.String{
MetaData: schematypes.MetaData{
Title: "WorkerType",
Description: `WorkerType to claim tasks for, combined with
'provisionerId' this identifies the pool of workers the machine
belongs to.`,
},
Pattern: `^[a-zA-Z0-9_-]{1,22}$`,
},
"workerGroup": schematypes.String{
MetaData: schematypes.MetaData{
Title: "WorkerGroup",
Description: `Group of workers this machine belongs to. This is any
identifier such that workerGroup and workerId uniquely identifies this
machine.`,
},
Pattern: `^[a-zA-Z0-9_-]{1,22}$`,
},
"workerId": schematypes.String{
MetaData: schematypes.MetaData{
Title: "WorkerId",
Description: `Identifier for this machine. This is any identifier such
that workerGroup and workerId uniquely identifies this machine.`,
},
Pattern: `^[a-zA-Z0-9_-]{1,22}$`,
},
"temporaryFolder": schematypes.String{
MetaData: schematypes.MetaData{
Title: "Temporary Folder",
Description: `Path to folder that can be used for temporary files and
folders, if folder doesn't exist it will be created, it will be
overwritten.`,
},
},
"logLevel": schematypes.StringEnum{
Options: []string{
logrus.DebugLevel.String(),
logrus.InfoLevel.String(),
logrus.WarnLevel.String(),
logrus.ErrorLevel.String(),
logrus.FatalLevel.String(),
logrus.PanicLevel.String(),
},
},
"serverIp": schematypes.String{},
"serverPort": schematypes.Integer{
Minimum: 0,
Maximum: 65535,
},
"networkInterface": schematypes.String{
MetaData: schematypes.MetaData{
Description: "Network device webhookserver should listen on. If not supplied, it binds to the interface from serverIp address",
},
},
"exposedPort": schematypes.Integer{
MetaData: schematypes.MetaData{
Description: "Port webhookserver should listen on. If not supplied, it uses the serverPort value.",
},
Minimum: 0,
Maximum: 65535,
},
"tlsCertificiate": schematypes.String{},
"tlsKey": schematypes.String{},
"statelessDNSSecret": schematypes.String{},
"statelessDNSDomain": schematypes.String{},
"maxLifeCycle": schematypes.Integer{
MetaData: schematypes.MetaData{
Title: "Max life cycle of worker",
Description: "Used to limit validity of hostname",
},
Minimum: 5 * 60,
Maximum: 31 * 24 * 60 * 60,
},
"minimumDiskSpace": schematypes.Integer{
MetaData: schematypes.MetaData{
Title: "Minimum Disk Space",
Description: `The minimum amount of disk space to have available
before starting on the next task. Garbage collector will do a
best-effort attempt at releasing resources to satisfy this limit`,
},
Minimum: 0,
Maximum: math.MaxInt64,
},
"minimumMemory": schematypes.Integer{
MetaData: schematypes.MetaData{
Title: "Minimum Memory",
Description: `The minimum amount of memory to have available
before starting on the next task. Garbage collector will do a
best-effort attempt at releasing resources to satisfy this limit`,
},
Minimum: 0,
Maximum: math.MaxInt64,
},
},
Required: []string{
"engine",
"engines",
"plugins",
"capacity",
"credentials",
"pollingInterval",
"reclaimOffset",
"provisionerId",
"workerType",
"workerGroup",
"workerId",
"temporaryFolder",
"logLevel",
"serverIp",
"serverPort",
"statelessDNSSecret",
"statelessDNSDomain",
"maxLifeCycle",
"minimumDiskSpace",
"minimumMemory",
},
}
}
func TestTaskManagerRunTask(t *testing.T) {
resolved := false
var serverURL string
var handler = func(w http.ResponseWriter, r *http.Request) {
if strings.Contains(r.URL.Path, "/artifacts/public/logs/live_backing.log") {
json.NewEncoder(w).Encode(&queue.S3ArtifactResponse{
PutURL: serverURL,
})
return
}
if strings.Contains(r.URL.Path, "/artifacts/public/logs/live.log") {
json.NewEncoder(w).Encode(&queue.RedirectArtifactResponse{})
return
}
if strings.Contains(r.URL.Path, "/task/abc/runs/1/completed") {
resolved = true
w.Header().Set("Content-Type", "application/json; charset=UTF-8")
json.NewEncoder(w).Encode(&queue.TaskStatusResponse{})
}
}
s := httptest.NewServer(http.HandlerFunc(handler))
serverURL = s.URL
defer s.Close()
tempPath := filepath.Join(os.TempDir(), slugid.Nice())
tempStorage, err := runtime.NewTemporaryStorage(tempPath)
if err != nil {
t.Fatal(err)
}
localServer, err := webhookserver.NewLocalServer(
[]byte{127, 0, 0, 1}, 60000,
"", 0,
"example.com",
"",
"",
"",
10*time.Minute,
)
if err != nil {
t.Error(err)
}
gc := &gc.GarbageCollector{}
environment := &runtime.Environment{
GarbageCollector: gc,
TemporaryStorage: tempStorage,
WebHookServer: localServer,
}
engineProvider := engines.Engines()["mock"]
engine, err := engineProvider.NewEngine(engines.EngineOptions{
Environment: environment,
Log: logger.WithField("engine", "mock"),
})
if err != nil {
t.Fatal(err.Error())
}
cfg := &configType{
QueueBaseURL: serverURL,
}
tm, err := newTaskManager(cfg, engine, MockPlugin{}, environment, logger.WithField("test", "TestTaskManagerRunTask"), gc)
if err != nil {
t.Fatal(err)
}
claim := &taskClaim{
taskID: "abc",
runID: 1,
taskClaim: &queue.TaskClaimResponse{
Credentials: struct {
AccessToken string `json:"accessToken"`
Certificate string `json:"certificate"`
ClientID string `json:"clientId"`
}{
AccessToken: "123",
ClientID: "abc",
Certificate: "",
},
TakenUntil: tcclient.Time(time.Now().Add(time.Minute * 5)),
},
definition: &queue.TaskDefinitionResponse{
Payload: []byte(`{"delay": 1,"function": "write-log","argument": "Hello World"}`),
},
}
tm.run(claim)
assert.True(t, resolved, "Task was not resolved")
}
func TestWorkerShutdown(t *testing.T) {
var resCount int32
var serverURL string
var handler = func(w http.ResponseWriter, r *http.Request) {
if strings.Contains(r.URL.Path, "/artifacts/public/logs/live_backing.log") {
json.NewEncoder(w).Encode(&queue.S3ArtifactResponse{
PutURL: serverURL,
})
return
}
if strings.Contains(r.URL.Path, "/artifacts/public/logs/live.log") {
json.NewEncoder(w).Encode(&queue.RedirectArtifactResponse{})
return
}
if strings.Contains(r.URL.Path, "exception") {
var exception queue.TaskExceptionRequest
err := json.NewDecoder(r.Body).Decode(&exception)
// Ignore errors for now
if err != nil {
return
}
assert.Equal(t, "worker-shutdown", exception.Reason)
atomic.AddInt32(&resCount, 1)
w.Header().Set("Content-Type", "application/json; charset=UTF-8")
json.NewEncoder(w).Encode(&queue.TaskStatusResponse{})
}
}
s := httptest.NewServer(http.HandlerFunc(handler))
serverURL = s.URL
defer s.Close()
tempPath := filepath.Join(os.TempDir(), slugid.Nice())
tempStorage, err := runtime.NewTemporaryStorage(tempPath)
if err != nil {
t.Fatal(err)
}
localServer, err := webhookserver.NewLocalServer(
[]byte{127, 0, 0, 1}, 60000,
"", 0,
"example.com",
"",
"",
"",
10*time.Minute,
)
if err != nil {
t.Error(err)
}
gc := &gc.GarbageCollector{}
environment := &runtime.Environment{
GarbageCollector: gc,
TemporaryStorage: tempStorage,
WebHookServer: localServer,
}
engineProvider := engines.Engines()["mock"]
engine, err := engineProvider.NewEngine(engines.EngineOptions{
Environment: environment,
Log: logger.WithField("engine", "mock"),
})
if err != nil {
t.Fatal(err.Error())
}
cfg := &configType{
QueueBaseURL: serverURL,
}
tm, err := newTaskManager(cfg, engine, MockPlugin{}, environment, logger.WithField("test", "TestRunTask"), gc)
if err != nil {
t.Fatal(err)
}
claims := []*taskClaim{
&taskClaim{
taskID: "abc",
runID: 1,
definition: &queue.TaskDefinitionResponse{
Payload: []byte(`{"delay": 5000,"function": "write-log","argument": "Hello World"}`),
},
taskClaim: &queue.TaskClaimResponse{
Credentials: struct {
AccessToken string `json:"accessToken"`
Certificate string `json:"certificate"`
ClientID string `json:"clientId"`
}{
AccessToken: "123",
ClientID: "abc",
Certificate: "",
},
TakenUntil: tcclient.Time(time.Now().Add(time.Minute * 5)),
},
},
&taskClaim{
taskID: "def",
runID: 0,
definition: &queue.TaskDefinitionResponse{
Payload: []byte(`{"delay": 5000,"function": "write-log","argument": "Hello World"}`),
},
taskClaim: &queue.TaskClaimResponse{
Credentials: struct {
AccessToken string `json:"accessToken"`
Certificate string `json:"certificate"`
ClientID string `json:"clientId"`
}{
AccessToken: "123",
ClientID: "abc",
Certificate: "",
},
TakenUntil: tcclient.Time(time.Now().Add(time.Minute * 5)),
},
},
}
var wg sync.WaitGroup
wg.Add(2)
go func() {
for _, c := range claims {
go func(claim *taskClaim) {
tm.run(claim)
wg.Done()
}(c)
}
}()
time.Sleep(500 * time.Millisecond)
assert.Equal(t, len(tm.RunningTasks()), 2)
close(tm.done)
tm.Stop()
wg.Wait()
assert.Equal(t, 0, len(tm.RunningTasks()))
assert.Equal(t, int32(2), atomic.LoadInt32(&resCount))
}
// Test is called to trigger a plugintest.Case to run
func (c Case) Test() {
runtimeEnvironment := newTestEnvironment()
testServer, err := webhookserver.NewTestServer()
nilOrPanic(err)
defer testServer.Stop()
runtimeEnvironment.WebHookServer = testServer
engineProvider := engines.Engines()["mock"]
engine, err := engineProvider.NewEngine(engines.EngineOptions{
Environment: runtimeEnvironment,
Log: runtimeEnvironment.Log.WithField("engine", "mock"),
// TODO: Add engine config
})
nilOrPanic(err, "engineProvider.NewEngine failed")
taskID := c.TaskID
if taskID == "" {
taskID = slugid.Nice()
}
context, controller, err := runtime.NewTaskContext(runtimeEnvironment.TemporaryStorage.NewFilePath(), runtime.TaskInfo{
TaskID: taskID,
RunID: c.RunID,
}, testServer)
nilOrPanic(err)
if c.QueueMock != nil {
controller.SetQueueClient(c.QueueMock)
}
sandboxBuilder, err := engine.NewSandboxBuilder(engines.SandboxOptions{
TaskContext: context,
Payload: parseEnginePayload(engine, c.Payload),
})
nilOrPanic(err, "engine.NewSandboxBuilder failed")
provider := plugins.Plugins()[c.Plugin]
assert(provider != nil, "Plugin does not exist! You tried to load: ", c.Plugin)
p, err := provider.NewPlugin(plugins.PluginOptions{
Environment: runtimeEnvironment,
Engine: engine,
Log: runtimeEnvironment.Log.WithField("plugin", c.Plugin),
Config: parsePluginConfig(provider, c.PluginConfig),
})
nilOrPanic(err, "pluginProvider.NewPlugin failed")
tp, err := p.NewTaskPlugin(plugins.TaskPluginOptions{
TaskInfo: &context.TaskInfo,
Payload: parsePluginPayload(p, c.Payload),
Log: runtimeEnvironment.Log.WithField("plugin", c.Plugin).WithField("taskId", taskID),
})
nilOrPanic(err, "plugin.NewTaskPlugin failed")
// taskPlugin can be nil, if the plugin doesn't want any hooks
if tp == nil {
tp = plugins.TaskPluginBase{}
}
options := Options{
Environment: runtimeEnvironment,
SandboxBuilder: sandboxBuilder,
Engine: engine,
Plugin: p,
TaskPlugin: tp,
}
err = tp.Prepare(context)
nilOrPanic(err, "taskPlugin.Prepare failed")
// Set environment variables and proxies
for key, val := range c.Env {
nilOrPanic(err, sandboxBuilder.SetEnvironmentVariable(key, val),
"Error setting env var: %s = %s", key, val)
}
for hostname, handler := range c.Proxies {
nilOrPanic(err, sandboxBuilder.AttachProxy(hostname, handler),
"Error attaching proxy for hostname: %s", hostname)
}
c.maybeRun(c.BeforeBuildSandbox, options)
err = tp.BuildSandbox(sandboxBuilder)
nilOrPanic(err, "taskPlugin.BuildSandbox failed")
c.maybeRun(c.AfterBuildSandbox, options)
c.maybeRun(c.BeforeStarted, options)
sandbox, err := sandboxBuilder.StartSandbox()
nilOrPanic(err, "sandboxBuilder.StartSandbox failed")
err = tp.Started(sandbox)
nilOrPanic(err, "taskPlugin.Started failed")
c.maybeRun(c.AfterStarted, options)
c.maybeRun(c.BeforeStopped, options)
resultSet, err := sandbox.WaitForResult()
nilOrPanic(err, "sandbox.WaitForResult failed")
assert(resultSet.Success() == c.EngineSuccess)
success, err := tp.Stopped(resultSet)
nilOrPanic(err, "taskPlugin.Stopped failed")
assert(success == c.PluginSuccess)
c.maybeRun(c.AfterStopped, options)
c.maybeRun(c.BeforeFinished, options)
controller.CloseLog()
err = tp.Finished(success)
nilOrPanic(err, "taskPlugin.Finished failed")
c.grepLog(context)
c.maybeRun(c.AfterFinished, options)
c.maybeRun(c.BeforeDisposed, options)
controller.Dispose()
err = tp.Dispose()
nilOrPanic(err, "taskPlugin.Dispose failed")
c.maybeRun(c.AfterDisposed, options)
}