diff --git a/docs/openapi.yaml b/docs/openapi.yaml index 9893ba2..39546e4 100644 --- a/docs/openapi.yaml +++ b/docs/openapi.yaml @@ -10,7 +10,7 @@ servers: info: description: API of the Fish node/cluster version: 1.0.0 - title: Aqurium Fish + title: Aquarium Fish contact: name: Sergei Parshev url: 'https://github.com/adobe/aquarium-fish' diff --git a/lib/drivers/aws/driver.go b/lib/drivers/aws/driver.go index d0dcaad..266a2fd 100644 --- a/lib/drivers/aws/driver.go +++ b/lib/drivers/aws/driver.go @@ -35,6 +35,21 @@ import ( "github.com/adobe/aquarium-fish/lib/util" ) +// Implements drivers.ResourceDriverFabric interface +type Fabric struct{} + +func (f *Fabric) Name() string { + return "aws" +} + +func (f *Fabric) NewResourceDriver() drivers.ResourceDriver { + return &Driver{} +} + +func init() { + drivers.FabricsList = append(drivers.FabricsList, &Fabric{}) +} + // Implements drivers.ResourceDriver interface type Driver struct { cfg Config @@ -47,10 +62,6 @@ type Driver struct { quotas_next_update time.Time } -func init() { - drivers.DriversList = append(drivers.DriversList, &Driver{}) -} - func (d *Driver) Name() string { return "aws" } diff --git a/lib/drivers/docker/driver.go b/lib/drivers/docker/driver.go index 14fa9b9..8b92fc3 100644 --- a/lib/drivers/docker/driver.go +++ b/lib/drivers/docker/driver.go @@ -30,6 +30,21 @@ import ( "github.com/adobe/aquarium-fish/lib/openapi/types" ) +// Implements drivers.ResourceDriverFabric interface +type Fabric struct{} + +func (f *Fabric) Name() string { + return "docker" +} + +func (f *Fabric) NewResourceDriver() drivers.ResourceDriver { + return &Driver{} +} + +func init() { + drivers.FabricsList = append(drivers.FabricsList, &Fabric{}) +} + // Implements drivers.ResourceDriver interface type Driver struct { cfg Config @@ -43,10 +58,6 @@ type Driver struct { docker_usage types.Resources // Used when the docker is remote } -func init() { - drivers.DriversList = append(drivers.DriversList, &Driver{}) -} - func (d *Driver) Name() string { return "docker" } diff --git a/lib/drivers/driver.go b/lib/drivers/driver.go index ec499c2..2eb98bb 100644 --- a/lib/drivers/driver.go +++ b/lib/drivers/driver.go @@ -21,7 +21,16 @@ const ( StatusAllocated = "ALLOCATED" ) -var DriversList []ResourceDriver +var FabricsList []ResourceDriverFabric + +// Fabric allows to generate new instances of the drivers +type ResourceDriverFabric interface { + // Name of the driver + Name() string + + // Generates new resource driver + NewResourceDriver() ResourceDriver +} type ResourceDriver interface { // Name of the driver diff --git a/lib/drivers/native/driver.go b/lib/drivers/native/driver.go index 2d10f95..6ea30f9 100644 --- a/lib/drivers/native/driver.go +++ b/lib/drivers/native/driver.go @@ -26,6 +26,21 @@ import ( "github.com/adobe/aquarium-fish/lib/openapi/types" ) +// Implements drivers.ResourceDriverFabric interface +type Fabric struct{} + +func (f *Fabric) Name() string { + return "native" +} + +func (f *Fabric) NewResourceDriver() drivers.ResourceDriver { + return &Driver{} +} + +func init() { + drivers.FabricsList = append(drivers.FabricsList, &Fabric{}) +} + // Implements drivers.ResourceDriver interface type Driver struct { cfg Config @@ -41,10 +56,6 @@ type EnvData struct { Disks map[string]string // Map with disk_name = mount_path } -func init() { - drivers.DriversList = append(drivers.DriversList, &Driver{}) -} - func (d *Driver) Name() string { return "native" } diff --git a/lib/drivers/test/driver.go b/lib/drivers/test/driver.go index ed224ae..7a1356b 100644 --- a/lib/drivers/test/driver.go +++ b/lib/drivers/test/driver.go @@ -27,6 +27,21 @@ import ( "github.com/adobe/aquarium-fish/lib/openapi/types" ) +// Implements drivers.ResourceDriverFabric interface +type Fabric struct{} + +func (f *Fabric) Name() string { + return "test" +} + +func (f *Fabric) NewResourceDriver() drivers.ResourceDriver { + return &Driver{} +} + +func init() { + drivers.FabricsList = append(drivers.FabricsList, &Fabric{}) +} + // Implements drivers.ResourceDriver interface type Driver struct { cfg Config @@ -34,10 +49,6 @@ type Driver struct { tasks_list []drivers.ResourceDriverTask } -func init() { - drivers.DriversList = append(drivers.DriversList, &Driver{}) -} - func (d *Driver) Name() string { return "test" } diff --git a/lib/drivers/vmx/driver.go b/lib/drivers/vmx/driver.go index b43706d..4d34220 100644 --- a/lib/drivers/vmx/driver.go +++ b/lib/drivers/vmx/driver.go @@ -30,6 +30,21 @@ import ( "github.com/adobe/aquarium-fish/lib/util" ) +// Implements drivers.ResourceDriverFabric interface +type Fabric struct{} + +func (f *Fabric) Name() string { + return "vmx" +} + +func (f *Fabric) NewResourceDriver() drivers.ResourceDriver { + return &Driver{} +} + +func init() { + drivers.FabricsList = append(drivers.FabricsList, &Fabric{}) +} + // Implements drivers.ResourceDriver interface type Driver struct { cfg Config @@ -40,10 +55,6 @@ type Driver struct { total_ram uint // In RAM GB } -func init() { - drivers.DriversList = append(drivers.DriversList, &Driver{}) -} - func (d *Driver) Name() string { return "vmx" } diff --git a/lib/fish/config.go b/lib/fish/config.go index b8029c6..c7b9693 100644 --- a/lib/fish/config.go +++ b/lib/fish/config.go @@ -39,7 +39,10 @@ type Config struct { DefaultResourceLifetime string `json:"default_resource_lifetime"` // Sets the lifetime of the resource which will be used if label definition one is not set - Drivers []ConfigDriver `json:"drivers"` // If specified - only the listed plugins will be loaded + // Configuration for the node drivers, if defined - only the listed plugins will be loaded + // Each configuration could instantinate the same driver multiple times by adding instance name + // separated from driver by slash symbol (like "/prod" - will create "prod" instance). + Drivers []ConfigDriver `json:"drivers"` } type ConfigDriver struct { diff --git a/lib/fish/drivers.go b/lib/fish/drivers.go index 244c255..ed23ab0 100644 --- a/lib/fish/drivers.go +++ b/lib/fish/drivers.go @@ -14,6 +14,7 @@ package fish import ( "fmt" + "strings" "github.com/adobe/aquarium-fish/lib/drivers" "github.com/adobe/aquarium-fish/lib/log" @@ -27,57 +28,56 @@ import ( _ "github.com/adobe/aquarium-fish/lib/drivers/test" ) -var drivers_enabled_list []drivers.ResourceDriver +var drivers_instances map[string]drivers.ResourceDriver func (f *Fish) DriverGet(name string) drivers.ResourceDriver { - for _, drv := range drivers_enabled_list { - if drv.Name() == name { - return drv - } + if drivers_instances == nil { + log.Error("Fish: Resource drivers are not initialized to request the driver instance:", name) + return nil } - return nil + drv, _ := drivers_instances[name] + return drv } +// Making the drivers instances map with specified names func (f *Fish) DriversSet() error { - var list []drivers.ResourceDriver + instances := make(map[string]drivers.ResourceDriver) - for _, drv := range drivers.DriversList { - en := false - if len(f.cfg.Drivers) == 0 { - // If no drivers is specified in the config - load all - en = true - } else { - for _, res := range f.cfg.Drivers { - if res.Name == drv.Name() { - en = true - break + if len(f.cfg.Drivers) == 0 { + // If no drivers instances are specified in the config - load all the drivers + for _, fbr := range drivers.FabricsList { + instances[fbr.Name()] = fbr.NewResourceDriver() + log.Info("Fish: Resource driver enabled:", fbr.Name()) + } + } else { + for _, fbr := range drivers.FabricsList { + // One driver could be used multiple times by config suffixes + for _, cfg := range f.cfg.Drivers { + log.Debug("Fish: Processing driver config:", cfg.Name, "vs", fbr.Name()) + if cfg.Name == fbr.Name() || strings.HasPrefix(cfg.Name, fbr.Name()+"/") { + instances[cfg.Name] = fbr.NewResourceDriver() + log.Info("Fish: Resource driver enabled:", fbr.Name(), "as", cfg.Name) } } } - if en { - log.Info("Fish: Resource driver enabled:", drv.Name()) - list = append(list, drv) - } else { - log.Info("Fish: Resource driver disabled:", drv.Name()) - } - } - if len(f.cfg.Drivers) > len(list) { - return fmt.Errorf("Unable to enable all the required drivers %s", f.cfg.Drivers) + if len(f.cfg.Drivers) > len(instances) { + return fmt.Errorf("Unable to enable all the required drivers %s", f.cfg.Drivers) + } } - drivers_enabled_list = list + drivers_instances = instances return nil } func (f *Fish) DriversPrepare(configs []ConfigDriver) (errs []error) { - not_skipped_drivers := drivers_enabled_list[:0] - for _, drv := range drivers_enabled_list { + activated_drivers_instances := make(map[string]drivers.ResourceDriver) + for name, drv := range drivers_instances { // Looking for the driver config var json_cfg []byte for _, cfg := range configs { - if drv.Name() == cfg.Name { + if name == cfg.Name { json_cfg = []byte(cfg.Cfg) break } @@ -87,12 +87,12 @@ func (f *Fish) DriversPrepare(configs []ConfigDriver) (errs []error) { errs = append(errs, err) log.Warn("Fish: Resource driver prepare failed:", drv.Name(), err) } else { - not_skipped_drivers = append(not_skipped_drivers, drv) + activated_drivers_instances[name] = drv log.Info("Fish: Resource driver activated:", drv.Name()) } } - drivers_enabled_list = not_skipped_drivers + drivers_instances = activated_drivers_instances return errs } diff --git a/tests/helpers.go b/tests/helpers.go deleted file mode 100644 index ece9e13..0000000 --- a/tests/helpers.go +++ /dev/null @@ -1,139 +0,0 @@ -/** - * Copyright 2021 Adobe. All rights reserved. - * This file is licensed to you under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. You may obtain a copy - * of the License at http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software distributed under - * the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS - * OF ANY KIND, either express or implied. See the License for the specific language - * governing permissions and limitations under the License. - */ - -package tests - -import ( - "bufio" - "context" - "fmt" - "os" - "os/exec" - "path/filepath" - "strings" - "testing" -) - -var fish_path = os.Getenv("FISH_PATH") // Full path to the aquarium-fish binary - -// Saves state of the running Aquarium Fish for particular test -type AFInstance struct { - workspace string - fishStop context.CancelFunc - running bool - - api_address string - admin_token string -} - -func RunAquariumFish(t *testing.T, cfg string) *AFInstance { - afi := &AFInstance{} - - afi.workspace = t.TempDir() - t.Log("INFO: Created workspace:", afi.workspace) - - os.WriteFile(filepath.Join(afi.workspace, "config.yml"), []byte(cfg), 0644) - t.Log("INFO: Stored config:", cfg) - - afi.fishStart(t) - - return afi -} - -// Will return url to access API of AquariumFish -func (afi *AFInstance) ApiAddress(path string) string { - return fmt.Sprintf("https://%s/%s", afi.api_address, path) -} - -// Returns admin token -func (afi *AFInstance) AdminToken() string { - return afi.admin_token -} - -// Check the fish instance is running -func (afi *AFInstance) IsRunning() bool { - return afi.running -} - -// Restart the application -func (afi *AFInstance) Restart(t *testing.T) error { - t.Log("INFO: Restarting:", afi.workspace) - afi.fishStop() - afi.fishStart(t) - return nil -} - -// Cleanup after the test execution -func (afi *AFInstance) Cleanup(t *testing.T) error { - t.Log("INFO: Cleaning up:", afi.workspace) - afi.fishStop() - os.RemoveAll(afi.workspace) - return nil -} - -func (afi *AFInstance) fishStart(t *testing.T) { - ctx, cancel := context.WithCancel(context.Background()) - afi.fishStop = cancel - - cmd := exec.CommandContext(ctx, fish_path, "-v", "debug", "-c", filepath.Join(afi.workspace, "config.yml")) - cmd.Dir = afi.workspace - r, _ := cmd.StdoutPipe() - cmd.Stderr = cmd.Stdout - - init_done := make(chan string) - scanner := bufio.NewScanner(r) - go func() { - // Listening for log and scan for token and address - for scanner.Scan() { - line := scanner.Text() - t.Log(line) - if strings.HasPrefix(line, "Admin user pass: ") { - val := strings.SplitN(strings.TrimSpace(line), "Admin user pass: ", 2) - if len(val) < 2 { - init_done <- "ERROR: No token after 'Admin user pass: '" - break - } - afi.admin_token = val[1] - } - if strings.Contains(line, "API listening on: ") { - val := strings.SplitN(strings.TrimSpace(line), "API listening on: ", 2) - if len(val) < 2 { - init_done <- "ERROR: No address after 'API listening on: '" - break - } - afi.api_address = val[1] - } - if strings.HasSuffix(line, "Fish initialized") { - // Found the needed values and continue to process to print the fish output for - // test debugging purposes - init_done <- "" - } - } - t.Log("Reading of AquariumFish output is done") - }() - - go func() { - afi.running = true - if err := cmd.Run(); err != nil { - t.Log("AquariumFish process was stopped:", err) - init_done <- fmt.Sprintf("ERROR: Fish was stopped with exit code: %v", err) - } - afi.running = false - r.Close() - }() - - failed := <-init_done - - if failed != "" { - t.Fatalf(failed) - } -} diff --git a/tests/multiple_driver_instances_test.go b/tests/multiple_driver_instances_test.go new file mode 100644 index 0000000..500ab9f --- /dev/null +++ b/tests/multiple_driver_instances_test.go @@ -0,0 +1,273 @@ +/** + * Copyright 2024 Adobe. All rights reserved. + * This file is licensed to you under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. You may obtain a copy + * of the License at http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software distributed under + * the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS + * OF ANY KIND, either express or implied. See the License for the specific language + * governing permissions and limitations under the License. + */ + +package tests + +import ( + "crypto/tls" + "fmt" + "net/http" + "testing" + "time" + + "github.com/google/uuid" + "github.com/steinfletcher/apitest" + + "github.com/adobe/aquarium-fish/lib/openapi/types" + h "github.com/adobe/aquarium-fish/tests/helper" +) + +// Check if driver instance with low limits will not run the Application and high limits will +// * Setup a number of drivers with different restrictions +// * Fail to Allocate Application with exceeding requirements +// * Success to Allocate application with fit requirements +func Test_multiple_driver_instances(t *testing.T) { + t.Parallel() + afi := h.NewAquariumFish(t, "node-1", `--- +node_location: test_loc + +api_address: 127.0.0.1:0 + +drivers: + - name: test/dev + cfg: + cpu_limit: 4 + ram_limit: 8 + - name: test/prod + cfg: + cpu_limit: 8 + ram_limit: 16`) + + t.Cleanup(func() { + afi.Cleanup(t) + }) + + defer func() { + if r := recover(); r != nil { + fmt.Println("Recovered in f", r) + } + }() + + tr := &http.Transport{ + TLSClientConfig: &tls.Config{InsecureSkipVerify: true}, + } + cli := &http.Client{ + Timeout: time.Second * 5, + Transport: tr, + } + + var label types.Label + t.Run("Create bad Label with test/dev driver", func(t *testing.T) { + apitest.New(). + EnableNetworking(cli). + Post(afi.ApiAddress("api/v1/label/")). + JSON(`{"name":"test-label", "version":1, "definitions": [{"driver":"test/dev", "resources":{"cpu":5,"ram":9}}]}`). + BasicAuth("admin", afi.AdminToken()). + Expect(t). + Status(http.StatusOK). + End(). + JSON(&label) + + if label.UID == uuid.Nil { + t.Fatalf("Label UID is incorrect: %v", label.UID) + } + }) + + var app types.Application + t.Run("Create Application", func(t *testing.T) { + apitest.New(). + EnableNetworking(cli). + Post(afi.ApiAddress("api/v1/application/")). + JSON(`{"label_UID":"`+label.UID.String()+`"}`). + BasicAuth("admin", afi.AdminToken()). + Expect(t). + Status(http.StatusOK). + End(). + JSON(&app) + + if app.UID == uuid.Nil { + t.Fatalf("Application UID is incorrect: %v", app.UID) + } + }) + + time.Sleep(10 * time.Second) + + var app_state types.ApplicationState + t.Run("Application should have state NEW in 10 sec", func(t *testing.T) { + apitest.New(). + EnableNetworking(cli). + Get(afi.ApiAddress("api/v1/application/"+app.UID.String()+"/state")). + BasicAuth("admin", afi.AdminToken()). + Expect(t). + Status(http.StatusOK). + End(). + JSON(&app_state) + + if app_state.Status != types.ApplicationStatusNEW { + t.Fatalf("Application Status is incorrect: %v", app_state.Status) + } + }) + + time.Sleep(10 * time.Second) + + t.Run("Application should have state NEW in 20 sec", func(t *testing.T) { + apitest.New(). + EnableNetworking(cli). + Get(afi.ApiAddress("api/v1/application/"+app.UID.String()+"/state")). + BasicAuth("admin", afi.AdminToken()). + Expect(t). + Status(http.StatusOK). + End(). + JSON(&app_state) + + if app_state.Status != types.ApplicationStatusNEW { + t.Fatalf("Application Status is incorrect: %v", app_state.Status) + } + }) + + time.Sleep(10 * time.Second) + + t.Run("Application should have state NEW in 30 sec", func(t *testing.T) { + apitest.New(). + EnableNetworking(cli). + Get(afi.ApiAddress("api/v1/application/"+app.UID.String()+"/state")). + BasicAuth("admin", afi.AdminToken()). + Expect(t). + Status(http.StatusOK). + End(). + JSON(&app_state) + + if app_state.Status != types.ApplicationStatusNEW { + t.Fatalf("Application Status is incorrect: %v", app_state.Status) + } + }) + + time.Sleep(10 * time.Second) + + t.Run("Application should have state NEW in 40 sec", func(t *testing.T) { + apitest.New(). + EnableNetworking(cli). + Get(afi.ApiAddress("api/v1/application/"+app.UID.String()+"/state")). + BasicAuth("admin", afi.AdminToken()). + Expect(t). + Status(http.StatusOK). + End(). + JSON(&app_state) + + if app_state.Status != types.ApplicationStatusNEW { + t.Fatalf("Application Status is incorrect: %v", app_state.Status) + } + }) + + t.Run("Deallocate the Application", func(t *testing.T) { + apitest.New(). + EnableNetworking(cli). + Get(afi.ApiAddress("api/v1/application/"+app.UID.String()+"/deallocate")). + BasicAuth("admin", afi.AdminToken()). + Expect(t). + Status(http.StatusOK). + End() + }) + + t.Run("Application should get RECALLED in 10 sec", func(t *testing.T) { + h.Retry(&h.Timer{Timeout: 5 * time.Second, Wait: 1 * time.Second}, t, func(r *h.R) { + apitest.New(). + EnableNetworking(cli). + Get(afi.ApiAddress("api/v1/application/"+app.UID.String()+"/state")). + BasicAuth("admin", afi.AdminToken()). + Expect(r). + Status(http.StatusOK). + End(). + JSON(&app_state) + + if app_state.Status != types.ApplicationStatusRECALLED { + r.Fatalf("Application Status is incorrect: %v", app_state.Status) + } + }) + }) + + t.Run("Create good Label with test/prod driver", func(t *testing.T) { + apitest.New(). + EnableNetworking(cli). + Post(afi.ApiAddress("api/v1/label/")). + JSON(`{"name":"test-label", "version":2, "definitions": [{"driver":"test/prod", "resources":{"cpu":5,"ram":9}}]}`). + BasicAuth("admin", afi.AdminToken()). + Expect(t). + Status(http.StatusOK). + End(). + JSON(&label) + + if label.UID == uuid.Nil { + t.Fatalf("Label UID is incorrect: %v", label.UID) + } + }) + + t.Run("Create Application", func(t *testing.T) { + apitest.New(). + EnableNetworking(cli). + Post(afi.ApiAddress("api/v1/application/")). + JSON(`{"label_UID":"`+label.UID.String()+`"}`). + BasicAuth("admin", afi.AdminToken()). + Expect(t). + Status(http.StatusOK). + End(). + JSON(&app) + + if app.UID == uuid.Nil { + t.Fatalf("Application UID is incorrect: %v", app.UID) + } + }) + + t.Run("Application should get ALLOCATED in 10 sec", func(t *testing.T) { + h.Retry(&h.Timer{Timeout: 10 * time.Second, Wait: 1 * time.Second}, t, func(r *h.R) { + apitest.New(). + EnableNetworking(cli). + Get(afi.ApiAddress("api/v1/application/"+app.UID.String()+"/state")). + BasicAuth("admin", afi.AdminToken()). + Expect(r). + Status(http.StatusOK). + End(). + JSON(&app_state) + + if app_state.Status != types.ApplicationStatusALLOCATED { + r.Fatalf("Application Status is incorrect: %v", app_state.Status) + } + }) + }) + + t.Run("Deallocate the Application", func(t *testing.T) { + apitest.New(). + EnableNetworking(cli). + Get(afi.ApiAddress("api/v1/application/"+app.UID.String()+"/deallocate")). + BasicAuth("admin", afi.AdminToken()). + Expect(t). + Status(http.StatusOK). + End() + }) + + t.Run("Application should get DEALLOCATED in 10 sec", func(t *testing.T) { + h.Retry(&h.Timer{Timeout: 10 * time.Second, Wait: 1 * time.Second}, t, func(r *h.R) { + apitest.New(). + EnableNetworking(cli). + Get(afi.ApiAddress("api/v1/application/"+app.UID.String()+"/state")). + BasicAuth("admin", afi.AdminToken()). + Expect(r). + Status(http.StatusOK). + End(). + JSON(&app_state) + + if app_state.Status != types.ApplicationStatusDEALLOCATED { + r.Fatalf("Application Status is incorrect: %v", app_state.Status) + } + }) + }) +} diff --git a/tests/retry.go b/tests/retry.go deleted file mode 100644 index 0777bbf..0000000 --- a/tests/retry.go +++ /dev/null @@ -1,195 +0,0 @@ -package tests - -import ( - "bytes" - "fmt" - "runtime" - "strings" - "time" -) - -// Failer is an interface compatible with testing.T. -type Failer interface { - Helper() - - // Log is called for the final test output - Log(args ...any) - - // FailNow is called when the retrying is abandoned. - FailNow() -} - -// R provides context for the retryer. -type R struct { - fail bool - done bool - output []string -} - -func (r *R) Helper() {} - -var runFailed = struct{}{} - -func (r *R) FailNow() { - r.fail = true - panic(runFailed) -} - -func (r *R) Fatal(args ...any) { - r.log(fmt.Sprint(args...)) - r.FailNow() -} - -func (r *R) Fatalf(format string, args ...any) { - r.log(fmt.Sprintf(format, args...)) - r.FailNow() -} - -func (r *R) Error(args ...any) { - r.log(fmt.Sprint(args...)) - r.fail = true -} - -func (r *R) Errorf(format string, args ...any) { - r.log(fmt.Sprintf(format, args...)) - r.fail = true -} - -func (r *R) Check(err error) { - if err != nil { - r.log(err.Error()) - r.FailNow() - } -} - -func (r *R) log(s string) { - r.output = append(r.output, decorate(s)) -} - -// Stop retrying, and fail the test with the specified error. -func (r *R) Stop(err error) { - r.log(err.Error()) - r.done = true -} - -func decorate(s string) string { - _, file, line, ok := runtime.Caller(3) - if ok { - n := strings.LastIndex(file, "/") - if n >= 0 { - file = file[n+1:] - } - } else { - file = "???" - line = 1 - } - return fmt.Sprintf("%s:%d: %s", file, line, s) -} - -func Retry(r Retryer, t Failer, f func(r *R)) { - t.Helper() - run(r, t, f) -} - -func dedup(a []string) string { - if len(a) == 0 { - return "" - } - seen := map[string]struct{}{} - var b bytes.Buffer - for _, s := range a { - if _, ok := seen[s]; ok { - continue - } - seen[s] = struct{}{} - b.WriteString(s) - b.WriteRune('\n') - } - return b.String() -} - -func run(r Retryer, t Failer, f func(r *R)) { - t.Helper() - rr := &R{} - - fail := func() { - t.Helper() - out := dedup(rr.output) - if out != "" { - t.Log(out) - } - t.FailNow() - } - - for r.Continue() { - func() { - defer func() { - if p := recover(); p != nil && p != runFailed { - panic(p) - } - }() - f(rr) - }() - - switch { - case rr.done: - fail() - return - case !rr.fail: - return - } - rr.fail = false - } - fail() -} - -// Retryer provides an interface for repeating operations -// until they succeed or an exit condition is met. -type Retryer interface { - // Continue returns true if the operation should be repeated, otherwise it - // returns false to indicate retrying should stop. - Continue() bool -} - -// Counter repeats an operation a given number of -// times and waits between subsequent operations. -type Counter struct { - Count int - Wait time.Duration - - count int -} - -func (r *Counter) Continue() bool { - if r.count == r.Count { - return false - } - if r.count > 0 { - time.Sleep(r.Wait) - } - r.count++ - return true -} - -// Timer repeats an operation for a given amount -// of time and waits between subsequent operations. -type Timer struct { - Timeout time.Duration - Wait time.Duration - - // stop is the timeout deadline. - // Set on the first invocation of Next(). - stop time.Time -} - -func (r *Timer) Continue() bool { - if r.stop.IsZero() { - r.stop = time.Now().Add(r.Timeout) - return true - } - if time.Now().After(r.stop) { - return false - } - time.Sleep(r.Wait) - return true -}