add caching for pricing and instance type data (#110)

Author: bwagner5

.github/workflows/ci.yaml

@@ -29,9 +29,6 @@ jobs:
 
     - name: Lints
       run: make spellcheck shellcheck
-
-    - name: Go Report Card Tests
-      run: make go-report-card-test
 
     - name: Brew Sync Dry run
       run: make homebrew-sync-dry-run

Dockerfile

@@ -1,4 +1,4 @@
-FROM public.ecr.aws/bitnami/golang:1.17 as builder
+FROM golang:1.17 as builder
 
 ## GOLANG env
 ARG GOPROXY="https://proxy.golang.org|direct"

Makefile

@@ -62,9 +62,6 @@ image:
 license-test:
 	${MAKEFILE_PATH}/test/license-test/run-license-test.sh
 
-go-report-card-test:
-	${MAKEFILE_PATH}/test/go-report-card-test/run-report-card-test.sh
-
 spellcheck:
 	${MAKEFILE_PATH}/test/readme-test/run-readme-spellcheck
 
@@ -110,7 +107,7 @@ build: compile
 
 release: build-binaries build-docker-images push-docker-images upload-resources-to-github
 
-test: spellcheck shellcheck unit-test license-test go-report-card-test e2e-test output-validation-test readme-codeblock-test
+test: spellcheck shellcheck unit-test license-test e2e-test output-validation-test readme-codeblock-test
 
 help:
 	@grep -E '^[a-zA-Z_-]+:.*$$' $(MAKEFILE_LIST) | sort

THIRD_PARTY_LICENSES

@@ -6,8 +6,8 @@ Copyright 2014-2015 Stripe, Inc.
 Copyright © 2015 Steve Francia <spf@spf13.com>
 ** go-yaml/yaml; version v2.2.2 -- https://github.com/go-yaml/yaml
 Copyright 2011-2016 Canonical Ltd.
-** go-ini/ini; version v1.57.0 -- https://github.com/go-ini/ini
-
+** gopkg.in/ini.v1; version v1.57.0 -- https://gopkg.in/ini.v1
+** gopkg.in/yaml.v2; version v2.3.0 -- https://gopkg.in/yaml.v2
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -309,6 +309,13 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 ** go-ghodss-yaml; version v1.0.0 -- https://github.com/ghodss/yaml
 Copyright (c) 2014 Sam Ghods
+** go-cache; version v2.1.0 -- https://github.com/patrickmn/go-cache
+Copyright (c) 2012-2019 Patrick Mylund Nielsen and the go-cache contributors
+** multierr; version v1.1.0 -- https://go.uber.org/multierr
+Copyright (c) 2017-2021 Uber Technologies, Inc.
+** atomic; version v1.4.0 -- https://go.uber.org/atomic
+Copyright (c) 2017-2021 Uber Technologies, Inc.
+
 
 The MIT License (MIT)
 

cmd/main.go

@@ -17,15 +17,20 @@ import (
 	"fmt"
 	"log"
 	"os"
+	"os/signal"
 	"strings"
 	"sync"
+	"syscall"
+	"time"
 
 	commandline "github.com/aws/amazon-ec2-instance-selector/v2/pkg/cli"
+	"github.com/aws/amazon-ec2-instance-selector/v2/pkg/env"
 	"github.com/aws/amazon-ec2-instance-selector/v2/pkg/selector"
 	"github.com/aws/amazon-ec2-instance-selector/v2/pkg/selector/outputs"
 	"github.com/aws/aws-sdk-go/aws/session"
 	homedir "github.com/mitchellh/go-homedir"
 	"github.com/spf13/cobra"
+	"go.uber.org/multierr"
 	"gopkg.in/ini.v1"
 )
 
@@ -35,6 +40,7 @@ const (
 	defaultRegionEnvVar = "AWS_DEFAULT_REGION"
 	defaultProfile      = "default"
 	awsConfigFile       = "~/.aws/config"
+	spotPricingDaysBack = 30
 
 	// cfnJSON is an output type
 	cfnJSON = "cfn-json"
@@ -90,6 +96,8 @@ const (
 	version    = "version"
 	region     = "region"
 	output     = "output"
+	cacheTTL   = "cache-ttl"
+	cacheDir   = "cache-dir"
 )
 
 var (
@@ -126,7 +134,7 @@ Full docs can be found at github.com/aws/amazon-` + binName
 	cli.IntMinMaxRangeFlags(vcpus, cli.StringMe("c"), nil, "Number of vcpus available to the instance type.")
 	cli.ByteQuantityMinMaxRangeFlags(memory, cli.StringMe("m"), nil, "Amount of Memory available (Example: 4 GiB)")
 	cli.RatioFlag(vcpusToMemoryRatio, nil, nil, "The ratio of vcpus to GiBs of memory. (Example: 1:2)")
-	cli.StringOptionsFlag(cpuArchitecture, cli.StringMe("a"), nil, "CPU architecture [x86_64/amd64, i386, or arm64]", []string{"x86_64", "amd64", "i386", "arm64"})
+	cli.StringOptionsFlag(cpuArchitecture, cli.StringMe("a"), nil, "CPU architecture [x86_64/amd64, x86_64_mac, i386, or arm64]", []string{"x86_64", "x86_64_mac", "amd64", "i386", "arm64"})
 	cli.IntMinMaxRangeFlags(gpus, cli.StringMe("g"), nil, "Total Number of GPUs (Example: 4)")
 	cli.ByteQuantityMinMaxRangeFlags(gpuMemoryTotal, nil, nil, "Number of GPUs' total memory (Example: 4 GiB)")
 	cli.StringOptionsFlag(placementGroupStrategy, nil, nil, "Placement group strategy: [cluster, partition, spread]", []string{"cluster", "partition", "spread"})
@@ -156,10 +164,12 @@ Full docs can be found at github.com/aws/amazon-` + binName
 
 	// Configuration Flags - These will be grouped at the bottom of the help flags
 
-	cli.ConfigIntFlag(maxResults, nil, cli.IntMe(20), "The maximum number of instance types that match your criteria to return")
+	cli.ConfigIntFlag(maxResults, nil, env.WithDefaultInt("EC2_INSTANCE_SELECTOR_MAX_RESULTS", 20), "The maximum number of instance types that match your criteria to return")
 	cli.ConfigStringFlag(profile, nil, nil, "AWS CLI profile to use for credentials and config", nil)
 	cli.ConfigStringFlag(region, cli.StringMe("r"), nil, "AWS Region to use for API requests (NOTE: if not passed in, uses AWS SDK default precedence)", nil)
 	cli.ConfigStringFlag(output, cli.StringMe("o"), nil, fmt.Sprintf("Specify the output format (%s)", strings.Join(cliOutputTypes, ", ")), nil)
+	cli.ConfigIntFlag(cacheTTL, nil, env.WithDefaultInt("EC2_INSTANCE_SELECTOR_CACHE_TTL", 168), "Cache TTLs in hours for pricing and instance type caches. Setting the cache to 0 will turn off caching and cleanup any on-disk caches.")
+	cli.ConfigPathFlag(cacheDir, nil, env.WithDefaultString("EC2_INSTANCE_SELECTOR_CACHE_DIR", "~/.ec2-instance-selector/"), "Directory to save the pricing and instance type caches")
 	cli.ConfigBoolFlag(verbose, cli.StringMe("v"), nil, "Verbose - will print out full instance specs")
 	cli.ConfigBoolFlag(help, cli.StringMe("h"), nil, "Help")
 	cli.ConfigBoolFlag(version, nil, nil, "Prints CLI version")
@@ -186,31 +196,36 @@ Full docs can be found at github.com/aws/amazon-` + binName
 		os.Exit(1)
 	}
 	flags[region] = sess.Config.Region
-
-	instanceSelector := selector.New(sess)
+	cacheTTLDuration := time.Hour * time.Duration(*cli.IntMe(flags[cacheTTL]))
+	instanceSelector := selector.NewWithCache(sess, cacheTTLDuration, *cli.StringMe(flags[cacheDir]))
+	shutdown := func() {
+		if err := instanceSelector.Save(); err != nil {
+			log.Printf("There was an error saving pricing caches: %v", err)
+		}
+	}
+	registerShutdown(shutdown)
 	outputFlag := cli.StringMe(flags[output])
 	if outputFlag != nil && *outputFlag == tableWideOutput {
 		// If output type is `table-wide`, simply print both prices for better comparison,
 		//   even if the actual filter is applied on any one of those based on usage class
-
-		// Save time by hydrating in parallel
-		wg := &sync.WaitGroup{}
-		wg.Add(2)
-		go func(waitGroup *sync.WaitGroup) {
-			defer waitGroup.Done()
-			_ = instanceSelector.EC2Pricing.HydrateOndemandCache()
-		}(wg)
-		go func(waitGroup *sync.WaitGroup) {
-			defer waitGroup.Done()
-			_ = instanceSelector.EC2Pricing.HydrateSpotCache(30)
-		}(wg)
-		wg.Wait()
+		// Save time by hydrating all caches in parallel
+		if err := hydrateCaches(*instanceSelector); err != nil {
+			log.Printf("%v", err)
+		}
 	} else if flags[pricePerHour] != nil {
 		// Else, if price filters are applied, only hydrate the respective cache as we don't have to print the prices
 		if flags[usageClass] == nil || *cli.StringMe(flags[usageClass]) == "on-demand" {
-			_ = instanceSelector.EC2Pricing.HydrateOndemandCache()
+			if instanceSelector.EC2Pricing.OnDemandCacheCount() == 0 {
+				if err := instanceSelector.EC2Pricing.RefreshOnDemandCache(); err != nil {
+					log.Printf("There was a problem refreshing the on-demand pricing cache: %v", err)
+				}
+			}
 		} else {
-			_ = instanceSelector.EC2Pricing.HydrateSpotCache(30)
+			if instanceSelector.EC2Pricing.SpotCacheCount() == 0 {
+				if err := instanceSelector.EC2Pricing.RefreshSpotCache(spotPricingDaysBack); err != nil {
+					log.Printf("There was a problem refreshing the spot pricing cache: %v", err)
+				}
+			}
 		}
 	}
 
@@ -290,6 +305,46 @@ Full docs can be found at github.com/aws/amazon-` + binName
 	if itemsTruncated > 0 {
 		log.Printf("%d entries were truncated, increase --%s to see more", itemsTruncated, maxResults)
 	}
+	shutdown()
+}
+
+func hydrateCaches(instanceSelector selector.Selector) (errs error) {
+	wg := &sync.WaitGroup{}
+	hydrateTasks := []func(*sync.WaitGroup) error{
+		func(waitGroup *sync.WaitGroup) error {
+			defer waitGroup.Done()
+			if instanceSelector.EC2Pricing.OnDemandCacheCount() == 0 {
+				if err := instanceSelector.EC2Pricing.RefreshOnDemandCache(); err != nil {
+					return multierr.Append(errs, fmt.Errorf("There was a problem refreshing the on-demand pricing cache: %w", err))
+				}
+			}
+			return nil
+		},
+		func(waitGroup *sync.WaitGroup) error {
+			defer waitGroup.Done()
+			if instanceSelector.EC2Pricing.SpotCacheCount() == 0 {
+				if err := instanceSelector.EC2Pricing.RefreshSpotCache(spotPricingDaysBack); err != nil {
+					return multierr.Append(errs, fmt.Errorf("There was a problem refreshing the spot pricing cache: %w", err))
+				}
+			}
+			return nil
+		},
+		func(waitGroup *sync.WaitGroup) error {
+			defer waitGroup.Done()
+			if instanceSelector.InstanceTypesProvider.CacheCount() == 0 {
+				if _, err := instanceSelector.InstanceTypesProvider.Get(nil); err != nil {
+					return multierr.Append(errs, fmt.Errorf("There was a problem refreshing the instance types cache: %w", err))
+				}
+			}
+			return nil
+		},
+	}
+	wg.Add(len(hydrateTasks))
+	for _, task := range hydrateTasks {
+		go task(wg)
+	}
+	wg.Wait()
+	return errs
 }
 
 func getOutputFn(outputFlag *string, currentFn selector.InstanceTypesOutputFn) selector.InstanceTypesOutputFn {
@@ -377,3 +432,12 @@ func getProfileRegion(profileName string) (string, error) {
 	}
 	return regionConfig.String(), nil
 }
+
+func registerShutdown(shutdown func()) {
+	sigs := make(chan os.Signal, 1)
+	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
+	go func() {
+		<-sigs
+		shutdown()
+	}()
+}

go.mod

@@ -9,6 +9,7 @@ require (
 	github.com/hashicorp/hcl v1.0.0
 	github.com/imdario/mergo v0.3.11
 	github.com/mitchellh/go-homedir v1.1.0
+	github.com/patrickmn/go-cache v2.1.0+incompatible
 	github.com/spf13/cobra v0.0.7
 	github.com/spf13/pflag v1.0.3
 	go.uber.org/multierr v1.1.0

go.sum

@@ -75,6 +75,8 @@ github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrk
 github.com/mitchellh/mapstructure v1.1.2/go.mod h1:FVVH3fgwuzCH5S8UJGiWEs2h04kUh9fWfEaFds41c1Y=
 github.com/mwitkow/go-conntrack v0.0.0-20161129095857-cc309e4a2223/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
 github.com/oklog/ulid v1.3.1/go.mod h1:CirwcVhetQ6Lv90oh/F+FBtV6XMibvdAFo93nm5qn4U=
+github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=
+github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
 github.com/pelletier/go-toml v1.2.0/go.mod h1:5z9KED0ma1S8pY6P1sdut58dfprrGBbd/94hg7ilaic=
 github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=

pkg/cli/flags.go

@@ -8,6 +8,7 @@ import (
 	"strings"
 
 	"github.com/aws/amazon-ec2-instance-selector/v2/pkg/bytequantity"
+	"github.com/mitchellh/go-homedir"
 	"github.com/spf13/pflag"
 )
 
@@ -90,6 +91,11 @@ func (cl *CommandLineInterface) RegexFlag(name string, shorthand *string, defaul
 	cl.RegexFlagOnFlagSet(cl.Command.Flags(), name, shorthand, defaultValue, description)
 }
 
+// PathFlag creates and registers a flag accepting a string representing a path and validates that it is a valid path.
+func (cl *CommandLineInterface) PathFlag(name string, shorthand *string, defaultValue *string, description string) {
+	cl.PathFlagOnFlagSet(cl.Command.Flags(), name, shorthand, defaultValue, description)
+}
+
 // StringOptionsFlag creates and registers a flag accepting a string and valid options for use in validation.
 func (cl *CommandLineInterface) StringOptionsFlag(name string, shorthand *string, defaultValue *string, description string, validOpts []string) {
 	cl.StringOptionsFlagOnFlagSet(cl.Command.Flags(), name, shorthand, defaultValue, description, validOpts)
@@ -107,11 +113,17 @@ func (cl *CommandLineInterface) ConfigStringFlag(name string, shorthand *string,
 }
 
 // ConfigStringSliceFlag creates and registers a flag accepting a list of strings.
-// Suite flags will be grouped in the middle of the output --help
+// Config flags will be grouped at the bottom in the output of --help
 func (cl *CommandLineInterface) ConfigStringSliceFlag(name string, shorthand *string, defaultValue []string, description string) {
 	cl.StringSliceFlagOnFlagSet(cl.Command.PersistentFlags(), name, shorthand, defaultValue, description)
 }
 
+// ConfigPathFlag creates and registers a flag accepting a string representing a path and validates that it is a valid path.
+// Config flags will be grouped at the bottom in the output of --help
+func (cl *CommandLineInterface) ConfigPathFlag(name string, shorthand *string, defaultValue *string, description string) {
+	cl.PathFlagOnFlagSet(cl.Command.PersistentFlags(), name, shorthand, defaultValue, description)
+}
+
 // ConfigIntFlag creates and registers a flag accepting an Integer for configuration purposes.
 // Config flags will be grouped at the bottom in the output of --help
 func (cl *CommandLineInterface) ConfigIntFlag(name string, shorthand *string, defaultValue *int, description string) {
@@ -378,3 +390,25 @@ func (cl *CommandLineInterface) RegexFlagOnFlagSet(flagSet *pflag.FlagSet, name
 	}
 	cl.StringFlagOnFlagSet(flagSet, name, shorthand, defaultValue, description, regexProcessor, regexValidator)
 }
+
+// PathFlagOnFlagSet creates and registers a flag accepting a string as a path
+func (cl *CommandLineInterface) PathFlagOnFlagSet(flagSet *pflag.FlagSet, name string, shorthand *string, defaultValue *string, description string) {
+	invalidInputMsg := fmt.Sprintf("Invalid path input for --%s. ", name)
+	pathProcessor := func(val interface{}) error {
+		if val == nil {
+			return nil
+		}
+		switch v := val.(type) {
+		case *string:
+			path, err := homedir.Expand(*v)
+			if err != nil {
+				return fmt.Errorf(invalidInputMsg + "Unable to expand path.")
+			}
+			cl.Flags[name] = &path
+		default:
+			return fmt.Errorf(invalidInputMsg + "Input type is unsupported.")
+		}
+		return nil
+	}
+	cl.StringFlagOnFlagSet(flagSet, name, shorthand, defaultValue, description, pathProcessor, nil)
+}