codevault/internal/utils/gshutdown/graceful_shutdown.go

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
}