bitmagnet/internal/queue/server/server.go

package server

import (
	"context"
	"database/sql"
	"errors"
	"sync"
	"time"

	"github.com/bitmagnet-io/bitmagnet/internal/database"
	"github.com/bitmagnet-io/bitmagnet/internal/database/dao"
	"github.com/bitmagnet-io/bitmagnet/internal/model"
	"github.com/bitmagnet-io/bitmagnet/internal/queue"
	"github.com/bitmagnet-io/bitmagnet/internal/queue/handler"
	"github.com/bitmagnet-io/bitmagnet/internal/workers/runner"
	"go.uber.org/zap"
	"gorm.io/gen"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

type server struct {
	daoProvider database.DaoTransactionProvider
	handlers    []handler.Handler
	gcInterval  time.Duration
	gcSemaphore chan struct{}
	draining    chan struct{}
	logger      *zap.Logger
}

func (s *server) Start(ctx context.Context, cancel context.CancelCauseFunc) (runner.Shutdowner, error) {
	handlers := make([]serverHandler, len(s.handlers))

	for i, h := range s.handlers {
		sh := serverHandler{
			Handler:     h,
			sem:         make(chan struct{}, h.Concurrency),
			draining:    s.draining,
			daoProvider: s.daoProvider,
			logger:      s.logger.With(zap.String("queue", h.Queue)),
		}

		handlers[i] = sh

		go sh.start(ctx)
	}

	go s.runGarbageCollection(ctx)

	return func(ctx context.Context) error {
		defer cancel(runner.ErrShutdownRequested)

		close(s.draining)

		var wg sync.WaitGroup

		wg.Add(len(handlers))

		for _, h := range handlers {
			go func(h serverHandler) {
				defer wg.Done()

				h.drain(ctx)
			}(h)
		}

		s.gcSemaphore <- struct{}{}

		wg.Wait()

		return nil
	}, nil
}

func (s *server) runGarbageCollection(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		case <-s.draining:
			return
		case s.gcSemaphore <- struct{}{}:
		}

		daoQ, err := s.daoProvider.Dao()
		if err != nil {
			s.logger.Error("error getting dao", zap.Error(err))
		}

		tx := daoQ.QueueJob.WithContext(ctx).Where(
			daoQ.QueueJob.Status.In(string(model.QueueJobStatusProcessed), string(model.QueueJobStatusFailed)),
		).
			UnderlyingDB().Where(
			"queue_jobs.ran_at + queue_jobs.archival_duration < ?::timestamptz",
			time.Now(),
		).Delete(&model.QueueJob{})

		if tx.Error != nil {
			s.logger.Error("error deleting old queue jobs", zap.Error(tx.Error))
		} else if tx.RowsAffected > 0 {
			s.logger.Debug("deleted old queue jobs", zap.Int64("count", tx.RowsAffected))
		}

		<-s.gcSemaphore

		select {
		case <-ctx.Done():
			return
		case <-s.draining:
			return
		case <-time.After(s.gcInterval):
			continue
		}
	}
}

type serverHandler struct {
	handler.Handler
	sem         chan struct{}
	draining    chan struct{}
	daoProvider database.DaoTransactionProvider
	// listenerChan chan pgconn.Notification
	logger *zap.Logger
}

func (h *serverHandler) start(ctx context.Context) {
	checkTicker := time.NewTicker(1)

	for {
		select {
		case <-ctx.Done():
			return
		case <-h.draining:
			return
		case <-checkTicker.C:
			select {
			case <-ctx.Done():
				return
			case <-h.draining:
				return
			case h.sem <- struct{}{}:
			}

			checkTicker.Reset(h.CheckInterval)

			go func() {
				jobID, _, err := h.handleJob(ctx)
				// if a job was found, we should check straight away for another job,
				// otherwise we wait for the check interval
				if err == nil && jobID != "" {
					checkTicker.Reset(1)
				}

				<-h.sem
			}()
		}
	}
}

func (h *serverHandler) handleJob(
	ctx context.Context,
	conds ...gen.Condition,
) (jobID string, processed bool, err error) {
	err = h.daoProvider.DaoTransaction(func(tx *dao.Query) error {
		job, findErr := tx.QueueJob.WithContext(ctx).Where(
			append(
				conds,
				tx.QueueJob.Queue.Eq(h.Queue),
				tx.QueueJob.Status.In(
					string(model.QueueJobStatusPending),
					string(model.QueueJobStatusRetry),
				),
				tx.QueueJob.RunAfter.Lte(time.Now()),
			)...,
		).Order(
			tx.QueueJob.Status.Eq(string(model.QueueJobStatusRetry)),
			tx.QueueJob.Priority,
			tx.QueueJob.RunAfter,
		).Clauses(clause.Locking{
			Strength: "UPDATE",
			Options:  "SKIP LOCKED",
		}).First()
		if findErr != nil {
			if errors.Is(findErr, gorm.ErrRecordNotFound) {
				return nil
			}

			return findErr
		}

		jobID = job.ID

		var jobErr error
		if job.Deadline.Valid && job.Deadline.Time.Before(time.Now()) {
			jobErr = ErrJobExceededDeadline

			h.logger.Debug("job deadline is in the past, skipping", zap.String("job_id", job.ID))
		} else {
			// check if the job is being retried and increment retry count accordingly
			if job.Status != model.QueueJobStatusPending {
				job.Retries++
			}

			jobErr = h.executeJob(ctx, *job)
		}

		job.RanAt = sql.NullTime{Time: time.Now(), Valid: true}

		if jobErr != nil {
			h.logger.Error("job failed", zap.Error(jobErr))

			if job.Retries < job.MaxRetries {
				job.Status = model.QueueJobStatusRetry
				job.RunAfter = queue.CalculateBackoff(job.Retries)
			} else {
				job.Status = model.QueueJobStatusFailed
			}

			job.Error = model.NewNullString(jobErr.Error())
		} else {
			job.Status = model.QueueJobStatusProcessed
			processed = true
		}

		_, updateErr := tx.QueueJob.WithContext(ctx).Updates(job)

		return updateErr
	})
	if err != nil {
		h.logger.Error("error handling job", zap.String("job_id", jobID), zap.Error(err))
	} else if processed {
		h.logger.Debug("job processed", zap.String("job_id", jobID))
	}

	return
}

func (h *serverHandler) executeJob(ctx context.Context, job model.QueueJob) error {
	run := h.Func(job)

	runCtx, runCancel := context.WithCancelCause(ctx)

	shutdowner, err := run(runCtx, runCancel)
	if err != nil {
		runCancel(err)
		return errors.Join(err, shutdowner(ctx))
	}

	select {
	case <-runCtx.Done():
		return extractContextCause(runCtx)
	case <-time.After(h.JobTimeout):
		runCancel(ErrJobExceededTimeout)
		return ErrJobExceededTimeout
	case <-h.draining:
		return shutdowner(runCtx)
	}
}

func extractContextCause(ctx context.Context) error {
	cause := context.Cause(ctx)

	if errors.Is(cause, runner.ErrCompleted) {
		return nil
	}

	return cause
}

func (h *serverHandler) drain(ctx context.Context) {
	for range h.Concurrency {
		select {
		case <-ctx.Done():
			return
		case h.sem <- struct{}{}:
		}
	}
}

var (
	ErrJobExceededDeadline = errors.New("the job did not complete before its deadline")
	ErrJobExceededTimeout  = errors.New("the job exceeded its timeout")
)