@@ -17,15 +17,20 @@ import (
1717 "fmt"
1818 "log"
1919 "os"
20+ "os/signal"
2021 "strings"
2122 "sync"
23+ "syscall"
24+ "time"
2225
2326 commandline "github.com/aws/amazon-ec2-instance-selector/v2/pkg/cli"
27+ "github.com/aws/amazon-ec2-instance-selector/v2/pkg/env"
2428 "github.com/aws/amazon-ec2-instance-selector/v2/pkg/selector"
2529 "github.com/aws/amazon-ec2-instance-selector/v2/pkg/selector/outputs"
2630 "github.com/aws/aws-sdk-go/aws/session"
2731 homedir "github.com/mitchellh/go-homedir"
2832 "github.com/spf13/cobra"
33+ "go.uber.org/multierr"
2934 "gopkg.in/ini.v1"
3035)
3136
@@ -35,6 +40,7 @@ const (
3540 defaultRegionEnvVar = "AWS_DEFAULT_REGION"
3641 defaultProfile = "default"
3742 awsConfigFile = "~/.aws/config"
43+ spotPricingDaysBack = 30
3844
3945 // cfnJSON is an output type
4046 cfnJSON = "cfn-json"
@@ -90,6 +96,8 @@ const (
9096 version = "version"
9197 region = "region"
9298 output = "output"
99+ cacheTTL = "cache-ttl"
100+ cacheDir = "cache-dir"
93101)
94102
95103var (
@@ -126,7 +134,7 @@ Full docs can be found at github.com/aws/amazon-` + binName
126134 cli .IntMinMaxRangeFlags (vcpus , cli .StringMe ("c" ), nil , "Number of vcpus available to the instance type." )
127135 cli .ByteQuantityMinMaxRangeFlags (memory , cli .StringMe ("m" ), nil , "Amount of Memory available (Example: 4 GiB)" )
128136 cli .RatioFlag (vcpusToMemoryRatio , nil , nil , "The ratio of vcpus to GiBs of memory. (Example: 1:2)" )
129- cli .StringOptionsFlag (cpuArchitecture , cli .StringMe ("a" ), nil , "CPU architecture [x86_64/amd64, i386, or arm64]" , []string {"x86_64" , "amd64" , "i386" , "arm64" })
137+ 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" })
130138 cli .IntMinMaxRangeFlags (gpus , cli .StringMe ("g" ), nil , "Total Number of GPUs (Example: 4)" )
131139 cli .ByteQuantityMinMaxRangeFlags (gpuMemoryTotal , nil , nil , "Number of GPUs' total memory (Example: 4 GiB)" )
132140 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
156164
157165 // Configuration Flags - These will be grouped at the bottom of the help flags
158166
159- cli .ConfigIntFlag (maxResults , nil , cli . IntMe ( 20 ), "The maximum number of instance types that match your criteria to return" )
167+ cli .ConfigIntFlag (maxResults , nil , env . WithDefaultInt ( "EC2_INSTANCE_SELECTOR_MAX_RESULTS" , 20 ), "The maximum number of instance types that match your criteria to return" )
160168 cli .ConfigStringFlag (profile , nil , nil , "AWS CLI profile to use for credentials and config" , nil )
161169 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 )
162170 cli .ConfigStringFlag (output , cli .StringMe ("o" ), nil , fmt .Sprintf ("Specify the output format (%s)" , strings .Join (cliOutputTypes , ", " )), nil )
171+ 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." )
172+ cli .ConfigPathFlag (cacheDir , nil , env .WithDefaultString ("EC2_INSTANCE_SELECTOR_CACHE_DIR" , "~/.ec2-instance-selector/" ), "Directory to save the pricing and instance type caches" )
163173 cli .ConfigBoolFlag (verbose , cli .StringMe ("v" ), nil , "Verbose - will print out full instance specs" )
164174 cli .ConfigBoolFlag (help , cli .StringMe ("h" ), nil , "Help" )
165175 cli .ConfigBoolFlag (version , nil , nil , "Prints CLI version" )
@@ -186,31 +196,36 @@ Full docs can be found at github.com/aws/amazon-` + binName
186196 os .Exit (1 )
187197 }
188198 flags [region ] = sess .Config .Region
189-
190- instanceSelector := selector .New (sess )
199+ cacheTTLDuration := time .Hour * time .Duration (* cli .IntMe (flags [cacheTTL ]))
200+ instanceSelector := selector .NewWithCache (sess , cacheTTLDuration , * cli .StringMe (flags [cacheDir ]))
201+ shutdown := func () {
202+ if err := instanceSelector .Save (); err != nil {
203+ log .Printf ("There was an error saving pricing caches: %v" , err )
204+ }
205+ }
206+ registerShutdown (shutdown )
191207 outputFlag := cli .StringMe (flags [output ])
192208 if outputFlag != nil && * outputFlag == tableWideOutput {
193209 // If output type is `table-wide`, simply print both prices for better comparison,
194210 // even if the actual filter is applied on any one of those based on usage class
195-
196- // Save time by hydrating in parallel
197- wg := & sync.WaitGroup {}
198- wg .Add (2 )
199- go func (waitGroup * sync.WaitGroup ) {
200- defer waitGroup .Done ()
201- _ = instanceSelector .EC2Pricing .HydrateOndemandCache ()
202- }(wg )
203- go func (waitGroup * sync.WaitGroup ) {
204- defer waitGroup .Done ()
205- _ = instanceSelector .EC2Pricing .HydrateSpotCache (30 )
206- }(wg )
207- wg .Wait ()
211+ // Save time by hydrating all caches in parallel
212+ if err := hydrateCaches (* instanceSelector ); err != nil {
213+ log .Printf ("%v" , err )
214+ }
208215 } else if flags [pricePerHour ] != nil {
209216 // Else, if price filters are applied, only hydrate the respective cache as we don't have to print the prices
210217 if flags [usageClass ] == nil || * cli .StringMe (flags [usageClass ]) == "on-demand" {
211- _ = instanceSelector .EC2Pricing .HydrateOndemandCache ()
218+ if instanceSelector .EC2Pricing .OnDemandCacheCount () == 0 {
219+ if err := instanceSelector .EC2Pricing .RefreshOnDemandCache (); err != nil {
220+ log .Printf ("There was a problem refreshing the on-demand pricing cache: %v" , err )
221+ }
222+ }
212223 } else {
213- _ = instanceSelector .EC2Pricing .HydrateSpotCache (30 )
224+ if instanceSelector .EC2Pricing .SpotCacheCount () == 0 {
225+ if err := instanceSelector .EC2Pricing .RefreshSpotCache (spotPricingDaysBack ); err != nil {
226+ log .Printf ("There was a problem refreshing the spot pricing cache: %v" , err )
227+ }
228+ }
214229 }
215230 }
216231
@@ -290,6 +305,46 @@ Full docs can be found at github.com/aws/amazon-` + binName
290305 if itemsTruncated > 0 {
291306 log .Printf ("%d entries were truncated, increase --%s to see more" , itemsTruncated , maxResults )
292307 }
308+ shutdown ()
309+ }
310+
311+ func hydrateCaches (instanceSelector selector.Selector ) (errs error ) {
312+ wg := & sync.WaitGroup {}
313+ hydrateTasks := []func (* sync.WaitGroup ) error {
314+ func (waitGroup * sync.WaitGroup ) error {
315+ defer waitGroup .Done ()
316+ if instanceSelector .EC2Pricing .OnDemandCacheCount () == 0 {
317+ if err := instanceSelector .EC2Pricing .RefreshOnDemandCache (); err != nil {
318+ return multierr .Append (errs , fmt .Errorf ("There was a problem refreshing the on-demand pricing cache: %w" , err ))
319+ }
320+ }
321+ return nil
322+ },
323+ func (waitGroup * sync.WaitGroup ) error {
324+ defer waitGroup .Done ()
325+ if instanceSelector .EC2Pricing .SpotCacheCount () == 0 {
326+ if err := instanceSelector .EC2Pricing .RefreshSpotCache (spotPricingDaysBack ); err != nil {
327+ return multierr .Append (errs , fmt .Errorf ("There was a problem refreshing the spot pricing cache: %w" , err ))
328+ }
329+ }
330+ return nil
331+ },
332+ func (waitGroup * sync.WaitGroup ) error {
333+ defer waitGroup .Done ()
334+ if instanceSelector .InstanceTypesProvider .CacheCount () == 0 {
335+ if _ , err := instanceSelector .InstanceTypesProvider .Get (nil ); err != nil {
336+ return multierr .Append (errs , fmt .Errorf ("There was a problem refreshing the instance types cache: %w" , err ))
337+ }
338+ }
339+ return nil
340+ },
341+ }
342+ wg .Add (len (hydrateTasks ))
343+ for _ , task := range hydrateTasks {
344+ go task (wg )
345+ }
346+ wg .Wait ()
347+ return errs
293348}
294349
295350func getOutputFn (outputFlag * string , currentFn selector.InstanceTypesOutputFn ) selector.InstanceTypesOutputFn {
@@ -377,3 +432,12 @@ func getProfileRegion(profileName string) (string, error) {
377432 }
378433 return regionConfig .String (), nil
379434}
435+
436+ func registerShutdown (shutdown func ()) {
437+ sigs := make (chan os.Signal , 1 )
438+ signal .Notify (sigs , syscall .SIGINT , syscall .SIGTERM )
439+ go func () {
440+ <- sigs
441+ shutdown ()
442+ }()
443+ }
0 commit comments