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) }