/
azbatch.go
94 lines (82 loc) · 3.43 KB
/
azbatch.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
/* (c) 2019, Blender Foundation
*
* Permission is hereby granted, free of charge, to any person obtaining
* a copy of this software and associated documentation files (the
* "Software"), to deal in the Software without restriction, including
* without limitation the rights to use, copy, modify, merge, publish,
* distribute, sublicense, and/or sell copies of the Software, and to
* permit persons to whom the Software is furnished to do so, subject to
* the following conditions:
*
* The above copyright notice and this permission notice shall be
* included in all copies or substantial portions of the Software.
*
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
* EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
* MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
* IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
* CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
* TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
* SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
*/
package azbatch
import (
"context"
"fmt"
"time"
"github.com/Azure/azure-sdk-for-go/services/batch/2018-12-01.8.0/batch"
"github.com/Azure/go-autorest/autorest/azure"
"github.com/Azure/go-autorest/autorest/date"
"github.com/sirupsen/logrus"
"github.com/Azure/flamenco-manager-azure/azauth"
"github.com/Azure/flamenco-manager-azure/azconfig"
"github.com/Azure/flamenco-manager-azure/aznetwork"
)
func getPoolClient(batchURL string) batch.PoolClient {
poolClient := batch.NewPoolClient(batchURL)
poolClient.Authorizer = azauth.Load(azure.PublicCloud.BatchManagementEndpoint)
// poolClient.RequestInspector = azdebug.LogRequest()
// poolClient.ResponseInspector = azdebug.LogResponse()
return poolClient
}
// CreatePool starts a pool of Flamenco Workers.
func CreatePool(config azconfig.AZConfig, netStack aznetwork.NetworkStack) {
ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(1*time.Minute))
defer cancel()
poolParams := PoolParameters(config, netStack)
batchURL := constructBatchURL(config)
createPoolIfNotExist(ctx, batchURL, poolParams)
}
func constructBatchURL(config azconfig.AZConfig) string {
return fmt.Sprintf("https://%s.%s.batch.azure.com", config.BatchAccountName, config.Location)
}
func createPoolIfNotExist(ctx context.Context, batchURL string, poolParams batch.PoolAddParameter) {
logger := logrus.WithField("pool_id", *poolParams.ID)
logrus.Info("fetching batch pools")
poolClient := getPoolClient(batchURL)
poolExists := false
resultPage, err := poolClient.List(ctx, "", "", "", nil, nil, nil, nil, nil)
if err != nil {
logrus.WithError(err).Fatal("unable to list existing pools")
}
for resultPage.NotDone() {
for _, foundPool := range resultPage.Values() {
logrus.WithField("found_id", *foundPool.ID).Info("found existing Azure Batch pool")
poolExists = poolExists || (*foundPool.ID == *poolParams.ID)
}
err := resultPage.NextWithContext(ctx)
if err != nil {
logrus.WithError(err).Fatal("unable to get next page of pools")
}
}
logger.WithField("pool_exists", poolExists).Debug("done listing pools")
if poolExists {
logger.Debug("Azure Batch pool exists")
return
}
_, err = poolClient.Add(ctx, poolParams, nil, nil, nil, &date.TimeRFC1123{Time: time.Now()})
if err != nil {
logger.WithError(err).Fatal("unable to add Azure Batch pool")
}
logger.Info("created Azure Batch pool")
}