extract util functions

acceptance_testing.go

@@ -1129,17 +1129,3 @@ func (a acceptanceTest) isEqualData(is *is.I, want, got Data) {
 		is.Equal(want.Bytes(), got.Bytes()) // data did not match (want != got)
 	}
 }
-
-func waitTimeout(wg *sync.WaitGroup, timeout time.Duration) bool {
-	c := make(chan struct{})
-	go func() {
-		defer close(c)
-		wg.Wait()
-	}()
-	select {
-	case <-c:
-		return false // completed normally
-	case <-time.After(timeout):
-		return true // timed out
-	}
-}

destination.go

@@ -179,7 +179,7 @@ func (a *destinationPluginAdapter) Run(ctx context.Context, stream cpluginv1.Des
 			if err == io.EOF {
 				// stream is closed
 				// wait for all acks to be sent back to Conduit
-				return a.waitForAcks(ctx)
+				return waitOrDone(&a.wgAckFuncs, ctx)
 			}
 			return fmt.Errorf("write stream error: %w", err)
 		}
@@ -235,25 +235,6 @@ func (a *destinationPluginAdapter) ackFunc(r Record, stream cpluginv1.Destinatio
 	}
 }
 
-func (a *destinationPluginAdapter) waitForAcks(ctx context.Context) error {
-	// wait for all acks to be sent back to Conduit
-	ackFuncsDone := make(chan struct{})
-	go func() {
-		a.wgAckFuncs.Wait()
-		close(ackFuncsDone)
-	}()
-	return a.waitForClose(ctx, ackFuncsDone)
-}
-
-func (a *destinationPluginAdapter) waitForClose(ctx context.Context, stop chan struct{}) error {
-	select {
-	case <-stop:
-		return nil
-	case <-ctx.Done():
-		return ctx.Err()
-	}
-}
-
 func (a *destinationPluginAdapter) Stop(ctx context.Context, req cpluginv1.DestinationStopRequest) (cpluginv1.DestinationStopResponse, error) {
 	// last thing we do is cancel context in Open
 	defer a.openCancel()
@@ -263,7 +244,7 @@ func (a *destinationPluginAdapter) Stop(ctx context.Context, req cpluginv1.Desti
 	defer cancel()
 
 	// wait for all acks to be sent back to Conduit
-	waitErr := a.waitForAcks(waitCtx)
+	waitErr := waitOrDone(&a.wgAckFuncs, waitCtx)
 	if waitErr != nil {
 		// just log error and continue to flush at least the processed records
 		Logger(ctx).Warn().Err(waitErr).Msg("failed to wait for all acks to be sent back to Conduit")
@@ -286,7 +267,7 @@ func (a *destinationPluginAdapter) Stop(ctx context.Context, req cpluginv1.Desti
 	// everything went as expected, let's cancel the context in Open and
 	// wait for Run to stop gracefully
 	a.openCancel()
-	err = a.waitForClose(ctx, a.runDone)
+	err = waitForClose(ctx, a.runDone)
 	return cpluginv1.DestinationStopResponse{}, err
 }
 

utils.go

@@ -0,0 +1,51 @@
+// Copyright © 2022 Meroxa, Inc.
+//
+// Licensed 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 CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package sdk
+
+import (
+	"context"
+	"sync"
+	"time"
+)
+
+// waitTimeout returns true if the given WaitGroup's counter is zero
+// before the given timeout is reached. Returns false otherwise.
+func waitTimeout(wg *sync.WaitGroup, timeout time.Duration) bool {
+	withTimeout, _ := context.WithTimeout(context.Background(), timeout)
+	return waitOrDone(wg, withTimeout) == nil
+}
+
+// waitTimeout returns nil if the given WaitGroup's counter is zero
+// before the given context is done. Returns the context's Err() otherwise.
+func waitOrDone(wg *sync.WaitGroup, ctx context.Context) error {
+	done := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(done)
+	}()
+	return waitForClose(ctx, done)
+}
+
+// waitForClose waits until the given channel receives a struct or until the given context is done.
+// If the channel receives a struct before the context is done, nil is returned.
+// Returns context's Err() otherwise.
+func waitForClose(ctx context.Context, stop chan struct{}) error {
+	select {
+	case <-stop:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}