Codevault initial commit

internal/utils/gshutdown/graceful_shutdown.go

@@ -0,0 +1,174 @@
+package gshutdown
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sync"
+	"time"
+
+	"go.uber.org/zap"
+)
+
+var (
+	ErrNotInitialized     = errors.New("seems that graceful shutdown has not been initialized")
+	ErrTimeout            = errors.New("timeout reached")
+	ErrContextFinished    = errors.New("context is finished, unable to add subtasks")
+	ErrContextNotFinished = errors.New("context is not finished, unable to await subtasks")
+)
+
+func CtxWithGracefulShutdown(ctx context.Context, logger *zap.Logger) context.Context {
+	if _, ok := ctx.Value(gracefulShutdownKey).(*gracefulShutdownData); ok {
+		return ctx
+	}
+	if logger == nil {
+		logger = zap.L().Named("graceful_shutdown")
+	}
+	data := gracefulShutdownData{
+		subtasks:       make(map[string]int),
+		operationsChan: make(chan operation, 1),
+		finishedChan:   make(chan struct{}),
+		logger:         logger,
+	}
+	ctx = context.WithValue(ctx, gracefulShutdownKey, &data)
+
+	go data.run(ctx)
+	data.logger.Info("initialized")
+
+	return ctx
+}
+
+func SubtaskStarted(ctx context.Context, subtaskName string) error {
+	if ctx.Err() != nil {
+		return ErrContextFinished
+	}
+	gs, ok := ctx.Value(gracefulShutdownKey).(*gracefulShutdownData)
+	if !ok {
+		return ErrNotInitialized
+	}
+	gs.operationsChan <- operation{Type: startOperationType, SubtaskName: subtaskName}
+	return nil
+}
+
+func SubtaskFinished(ctx context.Context, subtaskName string) error {
+	gs, ok := ctx.Value(gracefulShutdownKey).(*gracefulShutdownData)
+	if !ok {
+		return ErrNotInitialized
+	}
+	gs.operationsChan <- operation{Type: finishOperationType, SubtaskName: subtaskName}
+	return nil
+}
+
+func AwaitWithTimeout(ctx context.Context, timeout time.Duration) error {
+	if ctx.Err() == nil {
+		return fmt.Errorf("context is not finished yet")
+	}
+	gs, ok := ctx.Value(gracefulShutdownKey).(*gracefulShutdownData)
+	if !ok {
+		return ErrNotInitialized
+	}
+	timer := time.NewTimer(timeout)
+	select {
+	case <-timer.C:
+		return ErrTimeout
+	case <-gs.finishedChan:
+		return gs.err
+	}
+}
+
+type gracefulShutdownKeyType int
+
+const gracefulShutdownKey gracefulShutdownKeyType = iota
+
+type gracefulShutdownData struct {
+	mu             sync.Mutex
+	subtasks       map[string]int
+	operationsChan chan operation
+	finishedChan   chan struct{}
+	err            error
+	logger         *zap.Logger
+}
+
+type operationType string
+
+var (
+	startOperationType  operationType = "start"
+	finishOperationType operationType = "finish"
+)
+
+type operation struct {
+	Type        operationType
+	SubtaskName string
+}
+
+func (data *gracefulShutdownData) run(ctx context.Context) {
+	for {
+		select {
+		case <-ctx.Done():
+			if data.isDone() {
+				close(data.finishedChan)
+				data.logger.Info("no tasks to be awaited")
+				return
+			}
+		case op := <-data.operationsChan:
+			switch op.Type {
+			case startOperationType:
+				data.acknowledgeSubtask(op.SubtaskName)
+			default:
+				data.finishSubtask(op.SubtaskName)
+			}
+			if data.isDone() {
+				close(data.finishedChan)
+				data.logger.Info("all tasks successfully awaited")
+				return
+			}
+			data.logger.Debug("awaiting in progress", zap.Strings("to_be_awaited", data.toBeAwaited()))
+		}
+	}
+}
+
+func (data *gracefulShutdownData) acknowledgeSubtask(subtaskName string) {
+	data.mu.Lock()
+	defer data.mu.Unlock()
+
+	data.subtasks[subtaskName]++
+	data.logger.Info(
+		"successful acknowledge",
+		zap.String("subtask", subtaskName),
+		zap.Int("total_count", data.subtasks[subtaskName]),
+	)
+}
+
+func (data *gracefulShutdownData) finishSubtask(subtaskName string) {
+	data.mu.Lock()
+	defer data.mu.Unlock()
+
+	data.subtasks[subtaskName]--
+	data.logger.Info(
+		"successful finish",
+		zap.String("subtask", subtaskName),
+		zap.Int("left_count", data.subtasks[subtaskName]),
+	)
+	if data.subtasks[subtaskName] == 0 {
+		delete(data.subtasks, subtaskName)
+	}
+}
+
+func (data *gracefulShutdownData) isDone() bool {
+	data.mu.Lock()
+	defer data.mu.Unlock()
+	return len(data.subtasks) == 0
+}
+
+func (data *gracefulShutdownData) toBeAwaited() []string {
+	data.mu.Lock()
+	defer data.mu.Unlock()
+
+	res := make([]string, 0, len(data.subtasks))
+	for subtaskName, count := range data.subtasks {
+		for range count {
+			res = append(res, subtaskName)
+		}
+	}
+	return res
+}