lib/dispatchcloud/driver.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: AGPL-3.0
   4
   5 package dispatchcloud
   6
   7 import (
   8         "fmt"
   9         "time"
  10
  11         "git.curoverse.com/arvados.git/lib/cloud"
  12         "git.curoverse.com/arvados.git/lib/cloud/azure"
  13         "git.curoverse.com/arvados.git/lib/cloud/ec2"
  14         "git.curoverse.com/arvados.git/sdk/go/arvados"
  15         "github.com/sirupsen/logrus"
  16         "golang.org/x/crypto/ssh"
  17 )
  18
  19 var drivers = map[string]cloud.Driver{
  20         "azure": azure.Driver,
  21         "ec2":   ec2.Driver,
  22 }
  23
  24 func newInstanceSet(cluster *arvados.Cluster, setID cloud.InstanceSetID, logger logrus.FieldLogger) (cloud.InstanceSet, error) {
  25         driver, ok := drivers[cluster.CloudVMs.Driver]
  26         if !ok {
  27                 return nil, fmt.Errorf("unsupported cloud driver %q", cluster.CloudVMs.Driver)
  28         }
  29         is, err := driver.InstanceSet(cluster.CloudVMs.DriverParameters, setID, logger)
  30         if maxops := cluster.CloudVMs.MaxCloudOpsPerSecond; maxops > 0 {
  31                 is = &rateLimitedInstanceSet{
  32                         InstanceSet: is,
  33                         ticker:      time.NewTicker(time.Second / time.Duration(maxops)),
  34                 }
  35         }
  36         return is, err
  37 }
  38
  39 type rateLimitedInstanceSet struct {
  40         cloud.InstanceSet
  41         ticker *time.Ticker
  42 }
  43
  44 func (is rateLimitedInstanceSet) Create(it arvados.InstanceType, image cloud.ImageID, tags cloud.InstanceTags, init cloud.InitCommand, pk ssh.PublicKey) (cloud.Instance, error) {
  45         <-is.ticker.C
  46         inst, err := is.InstanceSet.Create(it, image, tags, init, pk)
  47         return &rateLimitedInstance{inst, is.ticker}, err
  48 }
  49
  50 type rateLimitedInstance struct {
  51         cloud.Instance
  52         ticker *time.Ticker
  53 }
  54
  55 func (inst *rateLimitedInstance) Destroy() error {
  56         <-inst.ticker.C
  57         return inst.Instance.Destroy()
  58 }