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
}
Example #3
0
// 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
}
Example #4
0
// 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)
}