1package worker
2
3import (
4	"context"
5	"errors"
6	"fmt"
7	"math/rand"
8	"time"
9
10	"code.cloudfoundry.org/lager"
11
12	"github.com/concourse/concourse/atc/db"
13)
14
15//go:generate counterfeiter . WorkerProvider
16
17type WorkerProvider interface {
18	RunningWorkers(lager.Logger) ([]Worker, error)
19
20	FindWorkerForContainer(
21		logger lager.Logger,
22		teamID int,
23		handle string,
24	) (Worker, bool, error)
25
26	FindWorkerForVolume(
27		logger lager.Logger,
28		teamID int,
29		handle string,
30	) (Worker, bool, error)
31
32	FindWorkersForContainerByOwner(
33		logger lager.Logger,
34		owner db.ContainerOwner,
35	) ([]Worker, error)
36
37	NewGardenWorker(
38		logger lager.Logger,
39		savedWorker db.Worker,
40		numBuildWorkers int,
41	) Worker
42}
43
44var (
45	ErrNoWorkers             = errors.New("no workers")
46	ErrFailedAcquirePoolLock = errors.New("failed to acquire pool lock")
47)
48
49type NoCompatibleWorkersError struct {
50	Spec WorkerSpec
51}
52
53func (err NoCompatibleWorkersError) Error() string {
54	return fmt.Sprintf("no workers satisfying: %s", err.Spec.Description())
55}
56
57//go:generate counterfeiter . Pool
58
59type Pool interface {
60	FindOrChooseWorker(
61		lager.Logger,
62		WorkerSpec,
63	) (Worker, error)
64
65	ContainerInWorker(
66		lager.Logger,
67		db.ContainerOwner,
68		WorkerSpec,
69	) (bool, error)
70
71	FindOrChooseWorkerForContainer(
72		context.Context,
73		lager.Logger,
74		db.ContainerOwner,
75		ContainerSpec,
76		WorkerSpec,
77		ContainerPlacementStrategy,
78	) (Worker, error)
79}
80
81type pool struct {
82	provider WorkerProvider
83	rand     *rand.Rand
84}
85
86func NewPool(
87	provider WorkerProvider,
88) Pool {
89	return &pool{
90		provider: provider,
91		rand:     rand.New(rand.NewSource(time.Now().UnixNano())),
92	}
93}
94
95func (pool *pool) allSatisfying(logger lager.Logger, spec WorkerSpec) ([]Worker, error) {
96	workers, err := pool.provider.RunningWorkers(logger)
97	if err != nil {
98		return nil, err
99	}
100
101	if len(workers) == 0 {
102		return nil, ErrNoWorkers
103	}
104
105	compatibleTeamWorkers := []Worker{}
106	compatibleGeneralWorkers := []Worker{}
107	for _, worker := range workers {
108		compatible := worker.Satisfies(logger, spec)
109		if compatible {
110			if worker.IsOwnedByTeam() {
111				compatibleTeamWorkers = append(compatibleTeamWorkers, worker)
112			} else {
113				compatibleGeneralWorkers = append(compatibleGeneralWorkers, worker)
114			}
115		}
116	}
117
118	if len(compatibleTeamWorkers) != 0 {
119		return compatibleTeamWorkers, nil
120	}
121
122	if len(compatibleGeneralWorkers) != 0 {
123		return compatibleGeneralWorkers, nil
124	}
125
126	return nil, NoCompatibleWorkersError{
127		Spec: spec,
128	}
129}
130
131func (pool *pool) ContainerInWorker(logger lager.Logger, owner db.ContainerOwner, workerSpec WorkerSpec) (bool, error) {
132	workersWithContainer, err := pool.provider.FindWorkersForContainerByOwner(
133		logger.Session("find-worker"),
134		owner,
135	)
136	if err != nil {
137		return false, err
138	}
139
140	compatibleWorkers, err := pool.allSatisfying(logger, workerSpec)
141	if err != nil {
142		return false, err
143	}
144
145	for _, w := range workersWithContainer {
146		for _, c := range compatibleWorkers {
147			if w.Name() == c.Name() {
148				return true, nil
149			}
150		}
151	}
152
153	return false, nil
154}
155
156func (pool *pool) FindOrChooseWorkerForContainer(
157	ctx context.Context,
158	logger lager.Logger,
159	owner db.ContainerOwner,
160	containerSpec ContainerSpec,
161	workerSpec WorkerSpec,
162	strategy ContainerPlacementStrategy,
163) (Worker, error) {
164	workersWithContainer, err := pool.provider.FindWorkersForContainerByOwner(
165		logger.Session("find-worker"),
166		owner,
167	)
168	if err != nil {
169		return nil, err
170	}
171
172	compatibleWorkers, err := pool.allSatisfying(logger, workerSpec)
173	if err != nil {
174		return nil, err
175	}
176
177	var worker Worker
178dance:
179	for _, w := range workersWithContainer {
180		for _, c := range compatibleWorkers {
181			if w.Name() == c.Name() {
182				worker = c
183				break dance
184			}
185		}
186	}
187
188	if worker == nil {
189		worker, err = strategy.Choose(logger, compatibleWorkers, containerSpec)
190		if err != nil {
191			return nil, err
192		}
193	}
194
195	return worker, nil
196}
197
198func (pool *pool) FindOrChooseWorker(
199	logger lager.Logger,
200	workerSpec WorkerSpec,
201) (Worker, error) {
202	workers, err := pool.allSatisfying(logger, workerSpec)
203	if err != nil {
204		return nil, err
205	}
206
207	return workers[rand.Intn(len(workers))], nil
208}
209