Implementing Timeout Feature on S3

aws/aws.go

@@ -21,9 +21,10 @@ import (
 
 // GetAllResources - Lists all aws resources
 func GetAllResources(c context.Context, query *Query, configObj config.Config) (*AwsAccountResources, error) {
-
 	configObj.AddExcludeAfterTime(query.ExcludeAfter)
 	configObj.AddIncludeAfterTime(query.IncludeAfter)
+	configObj.AddTimeout(query.Timeout)
+
 	configObj.KMSCustomerKeys.IncludeUnaliasedKeys = query.ListUnaliasedKMSKeys
 	account := AwsAccountResources{
 		Resources: make(map[string]AwsResources),
@@ -42,11 +43,21 @@ func GetAllResources(c context.Context, query *Query, configObj config.Config) (
 		registeredResources := GetAndInitRegisteredResources(cloudNukeSession, region)
 		for _, resource := range registeredResources {
 			if IsNukeable((*resource).ResourceName(), query.ResourceTypes) {
+
+				// PrepareContext sets up the resource context for execution, utilizing the context 'c' and the resource individual configuration.
+				// This function should be called after configuring the timeout to ensure proper execution context.
+				resourceConfig := (*resource).GetAndSetResourceConfig(configObj)
+				err := (*resource).PrepareContext(c, resourceConfig)
+				if err != nil {
+					return nil, err
+				}
+
 				spinner.UpdateText(
 					fmt.Sprintf("Searching %s resources in %s", (*resource).ResourceName(), region))
 				start := time.Now()
 				identifiers, err := (*resource).GetAndSetIdentifiers(c, configObj)
 				if err != nil {
+					logging.Errorf("Unable to retrieve %v, %v", (*resource).ResourceName(), err)
 					ge := report.GeneralError{
 						Error:        err,
 						Description:  fmt.Sprintf("Unable to retrieve %s", (*resource).ResourceName()),

aws/query.go

@@ -13,10 +13,11 @@ type Query struct {
 	ExcludeAfter         *time.Time
 	IncludeAfter         *time.Time
 	ListUnaliasedKMSKeys bool
+	Timeout              *time.Duration
 }
 
 // NewQuery configures and returns a Query struct that can be passed into the InspectResources method
-func NewQuery(regions, excludeRegions, resourceTypes, excludeResourceTypes []string, excludeAfter, includeAfter *time.Time, listUnaliasedKMSKeys bool) (*Query, error) {
+func NewQuery(regions, excludeRegions, resourceTypes, excludeResourceTypes []string, excludeAfter, includeAfter *time.Time, listUnaliasedKMSKeys bool, timeout *time.Duration) (*Query, error) {
 	q := &Query{
 		Regions:              regions,
 		ExcludeRegions:       excludeRegions,
@@ -25,6 +26,7 @@ func NewQuery(regions, excludeRegions, resourceTypes, excludeResourceTypes []str
 		ExcludeAfter:         excludeAfter,
 		IncludeAfter:         includeAfter,
 		ListUnaliasedKMSKeys: listUnaliasedKMSKeys,
+		Timeout:              timeout,
 	}
 
 	validationErr := q.Validate()

aws/query_test.go

@@ -1,10 +1,11 @@
 package aws
 
 import (
-	"github.com/aws/aws-sdk-go/aws"
 	"testing"
 	"time"
 
+	"github.com/aws/aws-sdk-go/aws"
+
 	"github.com/gruntwork-io/cloud-nuke/telemetry"
 
 	"github.com/stretchr/testify/require"
@@ -43,7 +44,7 @@ func TestNewQueryAcceptsValidExcludeAfterEntries(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.Name, func(t *testing.T) {
-			_, err := NewQuery(tc.Regions, tc.ExcludeRegions, tc.ResourceTypes, tc.ExcludeResourceTypes, tc.ExcludeAfter, tc.IncludeAfter, false)
+			_, err := NewQuery(tc.Regions, tc.ExcludeRegions, tc.ResourceTypes, tc.ExcludeResourceTypes, tc.ExcludeAfter, tc.IncludeAfter, false, nil)
 			require.NoError(t, err)
 		})
 	}

aws/resource.go

@@ -17,6 +17,9 @@ type AwsResource interface {
 	Nuke(identifiers []string) error
 	GetAndSetIdentifiers(c context.Context, configObj config.Config) ([]string, error)
 	IsNukable(string) (bool, error)
+
+	PrepareContext(context.Context, config.ResourceType) error
+	GetAndSetResourceConfig(config.Config) config.ResourceType
 }
 
 // AwsResources is a struct to hold multiple instances of AwsResource.

aws/resources/base_resource.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"time"
 
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/gruntwork-io/cloud-nuke/config"
@@ -16,52 +17,75 @@ import (
 type BaseAwsResource struct {
 	// A key-value of identifiers and nukable status
 	Nukables map[string]error
+	Timeout  time.Duration
+	Context  context.Context
+	cancel   context.CancelFunc
 }
 
-func (umpl *BaseAwsResource) Init(_ *session.Session) {
-	umpl.Nukables = make(map[string]error)
+func (br *BaseAwsResource) Init(_ *session.Session) {
+	br.Nukables = make(map[string]error)
 }
-func (umpl *BaseAwsResource) ResourceName() string {
+func (br *BaseAwsResource) ResourceName() string {
 	return "not implemented: ResourceName"
 }
-func (umpl *BaseAwsResource) ResourceIdentifiers() []string {
+func (br *BaseAwsResource) ResourceIdentifiers() []string {
 	return nil
 }
-func (umpl *BaseAwsResource) MaxBatchSize() int {
+func (br *BaseAwsResource) MaxBatchSize() int {
 	return 0
 }
-func (umpl *BaseAwsResource) Nuke(_ []string) error {
+func (br *BaseAwsResource) Nuke(_ []string) error {
 	return errors.New("not implemented: Nuke")
 }
-func (umpl *BaseAwsResource) GetAndSetIdentifiers(_ context.Context, _ config.Config) ([]string, error) {
+func (br *BaseAwsResource) GetAndSetIdentifiers(_ context.Context, _ config.Config) ([]string, error) {
 	return nil, errors.New("not implemented: GetAndSetIdentifiers")
 }
 
-func (umpl *BaseAwsResource) GetNukableStatus(identifier string) (error, bool) {
-	val, ok := umpl.Nukables[identifier]
+func (br *BaseAwsResource) GetNukableStatus(identifier string) (error, bool) {
+	val, ok := br.Nukables[identifier]
 	return val, ok
 }
 
-func (umpl *BaseAwsResource) SetNukableStatus(identifier string, err error) {
-	umpl.Nukables[identifier] = err
+func (br *BaseAwsResource) SetNukableStatus(identifier string, err error) {
+	br.Nukables[identifier] = err
+}
+func (br *BaseAwsResource) GetAndSetResourceConfig(_ config.Config) config.ResourceType {
+	return config.ResourceType{
+		Timeout: "",
+	}
+}
+
+func (br *BaseAwsResource) PrepareContext(parentContext context.Context, resourceConfig config.ResourceType) error {
+	if resourceConfig.Timeout == "" {
+		br.Context = parentContext
+		return nil
+	}
+
+	duration, err := time.ParseDuration(resourceConfig.Timeout)
+	if err != nil {
+		return err
+	}
+
+	br.Context, _ = context.WithTimeout(parentContext, duration)
+	return nil
 }
 
 // VerifyNukablePermissions performs nukable permission verification for each ID. For each ID, the function is
 // executed, and the result (error or success) is recorded using the SetNukableStatus method, indicating whether
 // the specified action is nukable
-func (umpl *BaseAwsResource) VerifyNukablePermissions(ids []*string, nukableCheckfn func(id *string) error) {
+func (br *BaseAwsResource) VerifyNukablePermissions(ids []*string, nukableCheckfn func(id *string) error) {
 	for _, id := range ids {
 		// skip if the id is already exists
-		if _, ok := umpl.GetNukableStatus(*id); ok {
+		if _, ok := br.GetNukableStatus(*id); ok {
 			continue
 		}
 		err := nukableCheckfn(id)
-		umpl.SetNukableStatus(*id, util.TransformAWSError(err))
+		br.SetNukableStatus(*id, util.TransformAWSError(err))
 	}
 }
 
-func (umpl *BaseAwsResource) IsNukable(identifier string) (bool, error) {
-	err, ok := umpl.Nukables[identifier]
+func (br *BaseAwsResource) IsNukable(identifier string) (bool, error) {
+	err, ok := br.Nukables[identifier]
 	if !ok {
 		return false, fmt.Errorf("-")
 	}