fix(wrapper): S3 Tables Wrapper is loading without waiting for the ad…

…dition process of deletedTablesCount
go-to-k · Jan 29, 2025 · 4e5110a · 4e5110a
1 parent c0307b2
commit 4e5110a
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/internal/wrapper/s3_tables_wrapper.go b/internal/wrapper/s3_tables_wrapper.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"strings"
+	"sync"
 	"sync/atomic"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
@@ -107,8 +108,10 @@ func (s *S3TablesWrapper) ClearBucket(
 
 	io.Logger.Info().Msgf("%v Checking...", bucketName)
 
-	progressCh := make(chan struct{})
 	var deletedTablesCount atomic.Int64
+	progressCh := make(chan struct{})
+	wg := sync.WaitGroup{}
+	wg.Add(1)
 	go func() {
 		for range progressCh {
 			count := deletedTablesCount.Add(1)
@@ -125,6 +128,7 @@ func (s *S3TablesWrapper) ClearBucket(
 		select {
 		case <-ctx.Done():
 			close(progressCh)
+			wg.Wait()
 			return &client.ClientError{
 				ResourceName: aws.String(bucketName),
 				Err:          ctx.Err(),
@@ -139,6 +143,7 @@ func (s *S3TablesWrapper) ClearBucket(
 		)
 		if err != nil {
 			close(progressCh)
+			wg.Wait()
 			return err
 		}
 		if len(output.Namespaces) == 0 {
@@ -149,6 +154,7 @@ func (s *S3TablesWrapper) ClearBucket(
 			for _, namespace := range summary.Namespace {
 				if err := sem.Acquire(ctx, 1); err != nil {
 					close(progressCh)
+					wg.Wait()
 					return err
 				}
 				eg.Go(func() error {
@@ -166,9 +172,11 @@ func (s *S3TablesWrapper) ClearBucket(
 
 	if err := eg.Wait(); err != nil {
 		close(progressCh)
+		wg.Wait()
 		return err
 	}
 	close(progressCh)
+	wg.Wait()
 
 	if !input.QuietMode {
 		if err := writer.Flush(); err != nil {

diff --git a/internal/wrapper/s3_tables_wrapper_test.go b/internal/wrapper/s3_tables_wrapper_test.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"reflect"
+	"sync"
 	"sync/atomic"
 	"testing"
 
@@ -659,14 +660,18 @@ func TestS3TablesWrapper_deleteNamespace(t *testing.T) {
 
 			progressCh := make(chan struct{})
 			var deletedCount atomic.Int64
+			wg := sync.WaitGroup{}
+			wg.Add(1)
 			go func() {
 				for range progressCh {
 					deletedCount.Add(1)
 				}
+				wg.Done()
 			}()
 
 			err := s3Tables.deleteNamespace(tt.args.ctx, tt.args.bucketArn, tt.args.bucketName, tt.args.namespace, progressCh)
 			close(progressCh)
+			wg.Wait()
 
 			if (err != nil) != tt.wantErr {
 				t.Errorf("error = %#v, wantErr %#v", err.Error(), tt.wantErr)