hardeepnarang10 · September 19, 2023 06:42 · Sep 19, 2023
diff --git a/graceful_stop_channels.go b/graceful_stop_channels.go
@@ -0,0 +1,43 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"time"
+)
+
+func gracefulStop(ctx context.Context, graceDuration time.Duration, stop func() error) error {
+	errChanIntermediary, errChanFinal := make(chan error), make(chan error)
+	defer close(errChanFinal)
+	defer close(errChanIntermediary)
+
+	go func(errChan chan<- error) {
+		if err := stop(); err != nil {
+			errChan <- fmt.Errorf("stop() call returned an error: %w", err)
+		}
+	}(errChanIntermediary)
+	<-ctx.Done()
+	preparedCtx, preparedCtxCancel := context.WithTimeout(context.Background(), graceDuration)
+	go func(prepCtx context.Context, prepCancel func(), errChannelIn <-chan error, errChannelOut chan<- error) {
+		for {
+			select {
+			case <-prepCtx.Done():
+				return
+			case err := <-errChannelIn:
+				errChannelOut <- err
+				prepCancel()
+				return
+			default:
+				continue
+			}
+		}
+	}(preparedCtx, preparedCtxCancel, errChanIntermediary, errChanFinal)
+	<-preparedCtx.Done()
+	if len(errChanFinal) != 0 {
+		if err := <-errChanFinal; err != nil {
+			log.Panic(err)
+		}
+	}
+	return nil
+}
diff --git a/graceful_stop_context.go b/graceful_stop_context.go
@@ -0,0 +1,32 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"time"
+)
+
+func gracefulStop(stop func() error, graceDuration time.Duration, graceTimeoutError error) error {
+	graceCtx, cancel := context.WithTimeout(context.Background(), graceDuration)
+	defer cancel()
+	errChan := make(chan error)
+	defer close(errChan)
+	context.AfterFunc(graceCtx, func() {
+		if err := stop(); err != nil {
+			errChan <- fmt.Errorf("stop() call returned an error: %w", err)
+		} else {
+			errChan <- nil
+		}
+	})
+	for {
+		select {
+		case <-graceCtx.Done():
+			return graceTimeoutError
+		case err := <-errChan:
+			if err != nil {
+				return err
+			}
+			return nil
+		}
+	}
+}