Migrate workers to kit/worker

Replace hand-rolled polling loops, semaphores, and WaitGroups
in all 7 background workers with go.gearno.de/kit/worker. Each
worker now implements Handler[T] (Claim/Process) and optionally
StaleRecoverer, gaining automatic Prometheus metrics and
OpenTelemetry tracing. Bumps kit from v0.3.0 to v0.5.0.

Signed-off-by: Bryan Frimin <bryan@getprobo.com>
This commit is contained in:
Bryan Frimin
2026-04-09 19:42:41 +02:00
parent 3ca0fac939
commit 43336078f4
13 changed files with 589 additions and 962 deletions

View File

@@ -1,123 +1,93 @@
# Go Worker
Background workers follow a poll-based pattern with bounded concurrency. The struct holds a `*pg.Client`, a `*log.Logger`, and tuning knobs (`interval`, `staleAfter`, `maxConcurrency`). Use functional options (`With*` functions) for the tuning knobs with sensible defaults.
Background workers use `go.gearno.de/kit/worker`. The kit handles the polling loop, semaphore-based concurrency, graceful shutdown, non-cancellable contexts, Prometheus metrics, and OpenTelemetry tracing. You only implement a **handler**.
## Run loop
## Handler interface
The `Run(ctx context.Context) error` method uses a `time.Ticker` in a `for`/`select` loop. On each tick it recovers stale rows, then drains available work via `processNext`. Work items are claimed inside a transaction with `FOR UPDATE SKIP LOCKED`, marked as processing, then handled concurrently in goroutines bounded by a semaphore channel. Use `context.WithoutCancel` for work that must complete even after shutdown, and `sync.WaitGroup` with `defer wg.Wait()` to ensure in-flight goroutines finish before `Run` returns.
Implement `worker.Handler[T]` with `Claim` and `Process`. Optionally implement `worker.StaleRecoverer` for stale row recovery.
```go
type (
FooWorker struct {
pg *pg.Client
logger *log.Logger
interval time.Duration
staleAfter time.Duration
maxConcurrency int
}
FooWorkerOption func(*FooWorker)
)
func NewFooWorker(
pgClient *pg.Client,
logger *log.Logger,
opts ...FooWorkerOption,
) *FooWorker {
w := &FooWorker{
pg: pgClient,
logger: logger,
interval: 10 * time.Second,
staleAfter: 5 * time.Minute,
maxConcurrency: 5,
}
for _, opt := range opts {
opt(w)
}
return w
type fooHandler struct {
pg *pg.Client
logger *log.Logger
staleAfter time.Duration
}
func (w *FooWorker) Run(ctx context.Context) error {
var (
wg sync.WaitGroup
sem = make(chan struct{}, w.maxConcurrency)
ticker = time.NewTicker(w.interval)
)
defer ticker.Stop()
defer wg.Wait()
func (h *fooHandler) Claim(ctx context.Context) (coredata.FooItem, error) {
var item coredata.FooItem
for {
select {
case <-ctx.Done():
return ctx.Err()
case <-ticker.C:
nonCancelableCtx := context.WithoutCancel(ctx)
w.recoverStaleRows(nonCancelableCtx)
for {
if err := w.processNext(ctx, sem, &wg); err != nil {
if !errors.Is(err, coredata.ErrResourceNotFound) {
w.logger.ErrorCtx(nonCancelableCtx, "cannot claim item", log.Error(err))
}
break
}
if err := h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
if err := item.LoadNextPendingForUpdateSkipLocked(ctx, tx); err != nil {
return err
}
now := time.Now()
item.Status = coredata.FooStatusProcessing
item.UpdatedAt = now
return item.Update(ctx, tx, coredata.NewNoScope())
},
); err != nil {
if errors.Is(err, coredata.ErrResourceNotFound) {
return coredata.FooItem{}, worker.ErrNoTask
}
return coredata.FooItem{}, err
}
return item, nil
}
func (h *fooHandler) Process(ctx context.Context, item coredata.FooItem) error {
if err := h.handle(ctx, &item); err != nil {
if failErr := h.fail(ctx, &item, err); failErr != nil {
h.logger.ErrorCtx(ctx, "cannot fail item", log.Error(failErr))
}
return err
}
return nil
}
func (h *fooHandler) RecoverStale(ctx context.Context) error {
return h.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) error {
return coredata.ResetStaleFooItems(ctx, conn, h.staleAfter)
},
)
}
```
## processNext
## Constructor
Claims one work item inside a transaction, marks it as processing, then handles it in a bounded goroutine:
The constructor builds the handler and returns `*worker.Worker[T]`. Use `worker.WithInterval` and `worker.WithMaxConcurrency` for tuning. Keep domain-specific options (e.g. timeouts, staleAfter) on the handler struct.
```go
func (w *FooWorker) processNext(ctx context.Context, sem chan struct{}, wg *sync.WaitGroup) error {
select {
case sem <- struct{}{}:
case <-ctx.Done():
return ctx.Err()
func NewFooWorker(
pgClient *pg.Client,
logger *log.Logger,
opts ...worker.Option,
) *worker.Worker[coredata.FooItem] {
h := &fooHandler{
pg: pgClient,
logger: logger,
staleAfter: 5 * time.Minute,
}
var (
item coredata.FooItem
now = time.Now()
nonCancelableCtx = context.WithoutCancel(ctx)
return worker.New(
"foo-worker",
h,
logger,
opts...,
)
if err := w.pg.WithTx(
nonCancelableCtx,
func(tx pg.Conn) error {
if err := item.LoadNextPendingForUpdateSkipLocked(nonCancelableCtx, tx); err != nil {
return err
}
item.Status = coredata.FooStatusProcessing
item.UpdatedAt = now
return item.Update(nonCancelableCtx, tx, coredata.NewNoScope())
},
); err != nil {
<-sem
return err
}
wg.Add(1)
go func(item coredata.FooItem) {
defer wg.Done()
defer func() { <-sem }()
if err := w.handle(nonCancelableCtx, &item); err != nil {
w.logger.ErrorCtx(nonCancelableCtx, "cannot process item", log.Error(err))
}
}(item)
return nil
}
```
## Key principles
- **Claim with `FOR UPDATE SKIP LOCKED`** — prevents multiple workers from picking the same row
- **Semaphore channel** — bounds goroutine concurrency to `maxConcurrency`
- **`context.WithoutCancel`** — in-flight work must complete even after shutdown
- **`defer wg.Wait()`** — `Run` blocks until all goroutines finish
- **Stale recovery** — on each tick, reset rows stuck in "processing" for longer than `staleAfter`
- **Drain loop** — keep calling `processNext` until no more pending items (`ErrResourceNotFound`)
- **Return `worker.ErrNoTask`** from `Claim` when no work is available (not the coredata sentinel)
- **Context is non-cancellable** — the kit provides `context.WithoutCancel` to both `Claim` and `Process`
- **Process handles its own failures** — update DB status on error, the kit only logs and records metrics
- **Stale recovery is optional** — implement `worker.StaleRecoverer` if the worker marks rows as "processing"
- **Kit provides observability** — Prometheus metrics (`worker_tasks_total`, `worker_task_duration_seconds`, etc.) and OTel traces are automatic

2
go.mod
View File

@@ -35,7 +35,7 @@ require (
github.com/vikstrous/dataloadgen v0.0.10
github.com/yuin/goldmark v1.4.13
go.gearno.de/crypto/uuid v0.1.1-0.20251208105319-3f587312a712
go.gearno.de/kit v0.3.0
go.gearno.de/kit v0.5.0
go.gearno.de/x/ref v0.0.0-20260216110753-a700c951377c
go.opentelemetry.io/otel v1.43.0
go.opentelemetry.io/otel/trace v1.43.0

4
go.sum
View File

@@ -332,8 +332,8 @@ github.com/yuin/goldmark v1.4.13 h1:fVcFKWvrslecOb/tg+Cc05dkeYx540o0FuFt3nUVDoE=
github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
go.gearno.de/crypto/uuid v0.1.1-0.20251208105319-3f587312a712 h1:J5ccbcxFuwxe6Oa9fVi9FqQOo+n17ni4wbl9t4NuEzc=
go.gearno.de/crypto/uuid v0.1.1-0.20251208105319-3f587312a712/go.mod h1:fnIIvKO9QnsyLO3ZJLJT3r8KZv/p0FOeT5eZKilYWXg=
go.gearno.de/kit v0.3.0 h1:c+0wY9ydGQIbLNvkcDv/1geC9htv1z2GDMOPxkZzoM4=
go.gearno.de/kit v0.3.0/go.mod h1:jWrI/mxd0F4GZApL0HgMextcEQoiy2YA1JVamSA/G0E=
go.gearno.de/kit v0.5.0 h1:rtxlR3LPi7Cq0SZurOuv+OgTVZ2/PSTG6PVtlMAndoM=
go.gearno.de/kit v0.5.0/go.mod h1:jWrI/mxd0F4GZApL0HgMextcEQoiy2YA1JVamSA/G0E=
go.gearno.de/x/panicf v0.1.1 h1:E3Cr9NB8Ry2EsvEG/1eHr7kplP3tEjTf5d56dTX64VQ=
go.gearno.de/x/panicf v0.1.1/go.mod h1:VnB8oF0UefMZcYeD4v+Wk4U5Z1uza7PHLlhT2CbNEbU=
go.gearno.de/x/ref v0.0.0-20260216110753-a700c951377c h1:rIVWwnNxHYu9aZhHkptXlNYTBJbY4ccaIAYjztVeaDc=

View File

@@ -21,6 +21,7 @@ import (
"go.gearno.de/kit/log"
"go.gearno.de/kit/pg"
"go.gearno.de/kit/worker"
"go.probo.inc/probo/pkg/connector"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/crypto/cipher"
@@ -35,16 +36,20 @@ type (
connectorRegistry *connector.ConnectorRegistry
logger *log.Logger
worker *SourceFetchWorker
sourceNameWorker *SourceNameWorker
fetchWorker *worker.Worker[coredata.AccessReviewCampaignSourceFetch]
sourceNameWorker *worker.Worker[coredata.AccessSource]
}
Option func(*Service)
Option func(*options)
options struct {
fetchInterval time.Duration
}
)
func WithFetchInterval(interval time.Duration) Option {
return func(s *Service) {
s.worker.interval = interval
return func(o *options) {
o.fetchInterval = interval
}
}
@@ -55,6 +60,11 @@ func NewService(
logger *log.Logger,
opts ...Option,
) *Service {
var o options
for _, opt := range opts {
opt(&o)
}
s := &Service{
pg: pgClient,
encryptionKey: encryptionKey,
@@ -62,7 +72,20 @@ func NewService(
logger: logger,
}
s.worker = NewSourceFetchWorker(s, pgClient, logger)
var fetchWorkerOpts []worker.Option
if o.fetchInterval > 0 {
fetchWorkerOpts = append(fetchWorkerOpts, worker.WithInterval(o.fetchInterval))
} else {
fetchWorkerOpts = append(fetchWorkerOpts, worker.WithInterval(30*time.Second))
}
fetchWorkerOpts = append(fetchWorkerOpts, worker.WithMaxConcurrency(20))
s.fetchWorker = NewSourceFetchWorker(
s,
pgClient,
logger,
fetchWorkerOpts...,
)
s.sourceNameWorker = NewSourceNameWorker(
pgClient,
encryptionKey,
@@ -70,10 +93,6 @@ func NewService(
logger.Named("source-name"),
)
for _, opt := range opts {
opt(s)
}
return s
}
@@ -134,14 +153,10 @@ func (s *Service) ResolveEntryOrganizationID(ctx context.Context, entryID gid.GI
}
func (s *Service) Run(ctx context.Context) error {
gCtx, cancel := context.WithCancel(context.WithoutCancel(ctx))
g, gCtx := errgroup.WithContext(gCtx)
g, gCtx := errgroup.WithContext(ctx)
g.Go(func() error { return s.worker.Run(gCtx) })
g.Go(func() error { return s.fetchWorker.Run(gCtx) })
g.Go(func() error { return s.sourceNameWorker.Run(gCtx) })
<-ctx.Done()
cancel()
return g.Wait()
}

View File

@@ -23,20 +23,20 @@ import (
"go.gearno.de/kit/log"
"go.gearno.de/kit/pg"
"go.gearno.de/kit/worker"
"go.probo.inc/probo/pkg/accessreview/drivers"
"go.probo.inc/probo/pkg/connector"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/crypto/cipher"
)
// SourceNameWorker polls for access sources that have a connector but no
// sourceNameHandler polls for access sources that have a connector but no
// synced name, resolves the provider instance name, and updates the source.
type SourceNameWorker struct {
type sourceNameHandler struct {
pg *pg.Client
encryptionKey cipher.EncryptionKey
connectorRegistry *connector.ConnectorRegistry
logger *log.Logger
interval time.Duration
}
func NewSourceNameWorker(
@@ -44,54 +44,49 @@ func NewSourceNameWorker(
encryptionKey cipher.EncryptionKey,
connectorRegistry *connector.ConnectorRegistry,
logger *log.Logger,
) *SourceNameWorker {
return &SourceNameWorker{
opts ...worker.Option,
) *worker.Worker[coredata.AccessSource] {
h := &sourceNameHandler{
pg: pgClient,
encryptionKey: encryptionKey,
connectorRegistry: connectorRegistry,
logger: logger,
interval: 10 * time.Second,
}
}
func (w *SourceNameWorker) Run(ctx context.Context) error {
w.logger.InfoCtx(ctx, "source name worker started",
log.String("interval", w.interval.String()),
defaultOpts := []worker.Option{
worker.WithInterval(10 * time.Second),
worker.WithMaxConcurrency(1),
}
return worker.New(
"source-name-worker",
h,
logger,
append(defaultOpts, opts...)...,
)
for {
select {
case <-ctx.Done():
w.logger.InfoCtx(context.WithoutCancel(ctx), "source name worker stopping")
return ctx.Err()
case <-time.After(w.interval):
nonCancelableCtx := context.WithoutCancel(ctx)
for {
if err := w.processNext(nonCancelableCtx); err != nil {
if !errors.Is(err, coredata.ErrNoAccessSourceNameSyncAvailable) {
w.logger.ErrorCtx(nonCancelableCtx, "cannot sync source name", log.Error(err))
}
break
}
}
}
}
}
func (w *SourceNameWorker) processNext(ctx context.Context) error {
func (h *sourceNameHandler) Claim(ctx context.Context) (coredata.AccessSource, error) {
var source coredata.AccessSource
err := w.pg.WithTx(
err := h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
return source.LoadNextUnsyncedNameForUpdateSkipLocked(ctx, tx)
},
)
if err != nil {
return err
if errors.Is(err, coredata.ErrNoAccessSourceNameSyncAvailable) {
return coredata.AccessSource{}, worker.ErrNoTask
}
return coredata.AccessSource{}, err
}
w.logger.InfoCtx(ctx, "syncing source name",
return source, nil
}
func (h *sourceNameHandler) Process(ctx context.Context, source coredata.AccessSource) error {
h.logger.InfoCtx(ctx, "syncing source name",
log.String("source_id", source.ID.String()),
log.String("current_name", source.Name),
)
@@ -101,7 +96,7 @@ func (w *SourceNameWorker) processNext(ctx context.Context) error {
resolver drivers.NameResolver
)
err = w.pg.WithTx(
err := h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
scope := coredata.NewScopeFromObjectID(source.ID)
@@ -109,7 +104,7 @@ func (w *SourceNameWorker) processNext(ctx context.Context) error {
return fmt.Errorf("source %s has no connector", source.ID)
}
if err := dbConnector.LoadByID(ctx, tx, scope, *source.ConnectorID, w.encryptionKey); err != nil {
if err := dbConnector.LoadByID(ctx, tx, scope, *source.ConnectorID, h.encryptionKey); err != nil {
return fmt.Errorf("cannot load connector %s: %w", *source.ConnectorID, err)
}
@@ -118,7 +113,7 @@ func (w *SourceNameWorker) processNext(ctx context.Context) error {
tokenBefore = oauth2Conn.AccessToken
}
httpClient, err := w.connectorHTTPClient(ctx, &dbConnector)
httpClient, err := h.connectorHTTPClient(ctx, &dbConnector)
if err != nil {
return fmt.Errorf("cannot create HTTP client for connector: %w", err)
}
@@ -126,18 +121,18 @@ func (w *SourceNameWorker) processNext(ctx context.Context) error {
if oauth2Conn, ok := dbConnector.Connection.(*connector.OAuth2Connection); ok {
if oauth2Conn.AccessToken != tokenBefore {
dbConnector.UpdatedAt = time.Now()
if err := dbConnector.Update(ctx, tx, scope, w.encryptionKey); err != nil {
if err := dbConnector.Update(ctx, tx, scope, h.encryptionKey); err != nil {
return fmt.Errorf("cannot persist refreshed token for connector %s: %w", *source.ConnectorID, err)
}
}
}
resolver = w.buildResolver(&dbConnector, httpClient)
resolver = h.buildResolver(&dbConnector, httpClient)
return nil
},
)
if err != nil {
w.logger.ErrorCtx(ctx, "cannot load connector for source name sync",
h.logger.ErrorCtx(ctx, "cannot load connector for source name sync",
log.String("source_id", source.ID.String()),
log.Error(err),
)
@@ -145,11 +140,11 @@ func (w *SourceNameWorker) processNext(ctx context.Context) error {
}
if resolver == nil {
w.logger.InfoCtx(ctx, "no name resolver for provider, keeping generic name",
h.logger.InfoCtx(ctx, "no name resolver for provider, keeping generic name",
log.String("source_id", source.ID.String()),
log.String("provider", dbConnector.Provider.String()),
)
return w.markNameSynced(ctx, &source)
return h.markNameSynced(ctx, &source)
}
resolveCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
@@ -157,7 +152,7 @@ func (w *SourceNameWorker) processNext(ctx context.Context) error {
instanceName, err := resolver.ResolveInstanceName(resolveCtx)
if err != nil {
w.logger.ErrorCtx(ctx, "cannot resolve instance name",
h.logger.ErrorCtx(ctx, "cannot resolve instance name",
log.String("source_id", source.ID.String()),
log.String("provider", dbConnector.Provider.String()),
log.Error(err),
@@ -166,31 +161,31 @@ func (w *SourceNameWorker) processNext(ctx context.Context) error {
}
if instanceName == "" {
w.logger.InfoCtx(ctx, "instance name is empty, keeping generic name",
h.logger.InfoCtx(ctx, "instance name is empty, keeping generic name",
log.String("source_id", source.ID.String()),
log.String("provider", dbConnector.Provider.String()),
)
return w.markNameSynced(ctx, &source)
return h.markNameSynced(ctx, &source)
}
displayName := drivers.ProviderDisplayName(dbConnector.Provider)
newName := displayName + " " + instanceName
w.logger.InfoCtx(ctx, "resolved source name",
h.logger.InfoCtx(ctx, "resolved source name",
log.String("source_id", source.ID.String()),
log.String("old_name", source.Name),
log.String("new_name", newName),
)
source.Name = newName
return w.markNameSynced(ctx, &source)
return h.markNameSynced(ctx, &source)
}
func (w *SourceNameWorker) markNameSynced(
func (h *sourceNameHandler) markNameSynced(
ctx context.Context,
source *coredata.AccessSource,
) error {
return w.pg.WithTx(
return h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
scope := coredata.NewScopeFromObjectID(source.ID)
@@ -212,7 +207,7 @@ func (w *SourceNameWorker) markNameSynced(
// For OAuth2 connections it uses RefreshableClient when a refresh config
// is registered for the provider, so that short-lived tokens are
// transparently refreshed.
func (w *SourceNameWorker) connectorHTTPClient(
func (h *sourceNameHandler) connectorHTTPClient(
ctx context.Context,
dbConnector *coredata.Connector,
) (*http.Client, error) {
@@ -221,8 +216,8 @@ func (w *SourceNameWorker) connectorHTTPClient(
return dbConnector.Connection.Client(ctx)
}
if w.connectorRegistry != nil {
refreshCfg := w.connectorRegistry.GetOAuth2RefreshConfig(string(dbConnector.Provider))
if h.connectorRegistry != nil {
refreshCfg := h.connectorRegistry.GetOAuth2RefreshConfig(string(dbConnector.Provider))
if refreshCfg != nil {
return oauth2Conn.RefreshableClient(ctx, *refreshCfg)
}
@@ -231,7 +226,7 @@ func (w *SourceNameWorker) connectorHTTPClient(
return oauth2Conn.Client(ctx)
}
func (w *SourceNameWorker) buildResolver(
func (h *sourceNameHandler) buildResolver(
dbConnector *coredata.Connector,
httpClient *http.Client,
) drivers.NameResolver {
@@ -249,7 +244,7 @@ func (w *SourceNameWorker) buildResolver(
case coredata.ConnectorProviderTally:
tallySettings, err := dbConnector.TallySettings()
if err != nil {
w.logger.Error("cannot read tally connector settings", log.Error(err))
h.logger.Error("cannot read tally connector settings", log.Error(err))
return nil
}
return drivers.NewTallyNameResolver(httpClient, tallySettings.OrganizationID)
@@ -262,21 +257,21 @@ func (w *SourceNameWorker) buildResolver(
case coredata.ConnectorProviderSentry:
sentrySettings, err := dbConnector.SentrySettings()
if err != nil {
w.logger.Error("cannot read sentry connector settings", log.Error(err))
h.logger.Error("cannot read sentry connector settings", log.Error(err))
return nil
}
return drivers.NewSentryNameResolver(httpClient, sentrySettings.OrganizationSlug)
case coredata.ConnectorProviderGitHub:
githubSettings, err := dbConnector.GitHubSettings()
if err != nil {
w.logger.Error("cannot read github connector settings", log.Error(err))
h.logger.Error("cannot read github connector settings", log.Error(err))
return nil
}
return drivers.NewGitHubNameResolver(httpClient, githubSettings.Organization)
case coredata.ConnectorProviderSupabase:
supabaseSettings, err := dbConnector.SupabaseSettings()
if err != nil {
w.logger.Error("cannot read supabase connector settings", log.Error(err))
h.logger.Error("cannot read supabase connector settings", log.Error(err))
return nil
}
return drivers.NewSupabaseNameResolver(supabaseSettings.OrganizationSlug)

View File

@@ -18,120 +18,54 @@ import (
"context"
"errors"
"fmt"
"sync"
"time"
"go.gearno.de/kit/log"
"go.gearno.de/kit/pg"
"go.gearno.de/kit/worker"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/gid"
)
type (
SourceFetchWorker struct {
svc *Service
pg *pg.Client
logger *log.Logger
interval time.Duration
staleAfter time.Duration
maxConcurrency int
}
SourceFetchWorkerOption func(*SourceFetchWorker)
)
func WithSourceFetchWorkerIntervalDuration(interval time.Duration) SourceFetchWorkerOption {
return func(w *SourceFetchWorker) {
w.interval = interval
}
}
func WithSourceFetchWorkerStaleAfter(staleAfter time.Duration) SourceFetchWorkerOption {
return func(w *SourceFetchWorker) {
w.staleAfter = staleAfter
}
}
func WithSourceFetchWorkerMaxConcurrency(maxConcurrency int) SourceFetchWorkerOption {
return func(w *SourceFetchWorker) {
w.maxConcurrency = maxConcurrency
}
type sourceFetchHandler struct {
svc *Service
pg *pg.Client
logger *log.Logger
staleAfter time.Duration
}
func NewSourceFetchWorker(
svc *Service,
pgClient *pg.Client,
logger *log.Logger,
opts ...SourceFetchWorkerOption,
) *SourceFetchWorker {
w := &SourceFetchWorker{
svc: svc,
pg: pgClient,
logger: logger,
interval: 30 * time.Second,
staleAfter: 5 * time.Minute,
maxConcurrency: 20,
opts ...worker.Option,
) *worker.Worker[coredata.AccessReviewCampaignSourceFetch] {
h := &sourceFetchHandler{
svc: svc,
pg: pgClient,
logger: logger,
staleAfter: 5 * time.Minute,
}
for _, opt := range opts {
opt(w)
}
return w
return worker.New(
"source-fetch-worker",
h,
logger,
opts...,
)
}
func (w *SourceFetchWorker) Run(ctx context.Context) error {
var (
wg sync.WaitGroup
sem = make(chan struct{}, w.maxConcurrency)
ticker = time.NewTicker(w.interval)
)
defer ticker.Stop()
defer wg.Wait()
func (h *sourceFetchHandler) Claim(ctx context.Context) (coredata.AccessReviewCampaignSourceFetch, error) {
var sourceFetch coredata.AccessReviewCampaignSourceFetch
for {
select {
case <-ctx.Done():
return ctx.Err()
case <-ticker.C:
nonCancelableCtx := context.WithoutCancel(ctx)
w.recoverStaleRows(nonCancelableCtx)
for {
if err := w.processNext(ctx, sem, &wg); err != nil {
if !errors.Is(err, coredata.ErrNoAccessReviewCampaignSourceFetchAvailable) {
w.logger.ErrorCtx(nonCancelableCtx, "cannot claim item", log.Error(err))
}
break
}
}
}
}
}
func (w *SourceFetchWorker) processNext(
ctx context.Context,
sem chan struct{},
wg *sync.WaitGroup,
) error {
select {
case sem <- struct{}{}:
case <-ctx.Done():
return ctx.Err()
}
var (
sourceFetch coredata.AccessReviewCampaignSourceFetch
now = time.Now()
nonCancelableCtx = context.WithoutCancel(ctx)
)
if err := w.pg.WithTx(
nonCancelableCtx,
if err := h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
if err := sourceFetch.LoadNextQueuedForUpdateSkipLocked(nonCancelableCtx, tx); err != nil {
return err // sentinel errors checked by caller
if err := sourceFetch.LoadNextQueuedForUpdateSkipLocked(ctx, tx); err != nil {
return err
}
now := time.Now()
sourceFetch.Status = coredata.AccessReviewCampaignSourceFetchStatusFetching
sourceFetch.AttemptCount++
sourceFetch.LastError = nil
@@ -140,38 +74,60 @@ func (w *SourceFetchWorker) processNext(
sourceFetch.UpdatedAt = now
scope := coredata.NewScope(sourceFetch.TenantID)
if err := sourceFetch.Update(nonCancelableCtx, tx, scope); err != nil {
if err := sourceFetch.Update(ctx, tx, scope); err != nil {
return fmt.Errorf("cannot update source fetch status: %w", err)
}
return nil
},
); err != nil {
<-sem
return fmt.Errorf("cannot claim source fetch: %w", err)
if errors.Is(err, coredata.ErrNoAccessReviewCampaignSourceFetchAvailable) {
return coredata.AccessReviewCampaignSourceFetch{}, worker.ErrNoTask
}
return coredata.AccessReviewCampaignSourceFetch{}, fmt.Errorf("cannot claim source fetch: %w", err)
}
wg.Add(1)
go func(sourceFetch coredata.AccessReviewCampaignSourceFetch) {
defer wg.Done()
defer func() { <-sem }()
if err := w.handle(nonCancelableCtx, &sourceFetch); err != nil {
w.logger.ErrorCtx(nonCancelableCtx, "cannot process source fetch", log.Error(err))
}
}(sourceFetch)
return nil
return sourceFetch, nil
}
func (w *SourceFetchWorker) handle(
func (h *sourceFetchHandler) Process(ctx context.Context, sourceFetch coredata.AccessReviewCampaignSourceFetch) error {
return h.handle(ctx, &sourceFetch)
}
func (h *sourceFetchHandler) RecoverStale(ctx context.Context) error {
now := time.Now()
staleThreshold := now.Add(-h.staleAfter)
return h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
var fetches coredata.AccessReviewCampaignSourceFetches
count, err := fetches.RecoverStale(ctx, tx, staleThreshold, now)
if err != nil {
return fmt.Errorf("cannot recover stale source fetches: %w", err)
}
if count > 0 {
h.logger.InfoCtx(
ctx,
"recovered stale source fetches",
log.Int64("count", count),
)
}
return nil
},
)
}
func (h *sourceFetchHandler) handle(
ctx context.Context,
sourceFetch *coredata.AccessReviewCampaignSourceFetch,
) error {
scope := coredata.NewScope(sourceFetch.TenantID)
campaign, err := w.svc.Campaigns(scope).Get(ctx, sourceFetch.AccessReviewCampaignID)
campaign, err := h.svc.Campaigns(scope).Get(ctx, sourceFetch.AccessReviewCampaignID)
if err != nil {
commitErr := w.commitFailedSourceFetch(
commitErr := h.commitFailedSourceFetch(
ctx,
sourceFetch,
fmt.Errorf("cannot load campaign: %w", err),
@@ -182,60 +138,31 @@ func (w *SourceFetchWorker) handle(
return fmt.Errorf("cannot load campaign: %w", err)
}
count, err := w.svc.Engine(scope).FetchSource(ctx, campaign, sourceFetch.AccessSourceID)
count, err := h.svc.Engine(scope).FetchSource(ctx, campaign, sourceFetch.AccessSourceID)
if err != nil {
commitErr := w.commitFailedSourceFetch(ctx, sourceFetch, err)
commitErr := h.commitFailedSourceFetch(ctx, sourceFetch, err)
if commitErr != nil {
return fmt.Errorf("cannot fetch source: %w, and cannot commit failed source fetch: %w", err, commitErr)
}
if finalizeErr := w.finalizeCampaignFetchLifecycle(ctx, sourceFetch.TenantID, sourceFetch.AccessReviewCampaignID); finalizeErr != nil {
if finalizeErr := h.finalizeCampaignFetchLifecycle(ctx, sourceFetch.TenantID, sourceFetch.AccessReviewCampaignID); finalizeErr != nil {
return fmt.Errorf("cannot finalize campaign after failed source fetch: %w", finalizeErr)
}
return fmt.Errorf("cannot fetch source: %w", err)
}
if err := w.commitSuccessfulSourceFetch(ctx, sourceFetch, count); err != nil {
if err := h.commitSuccessfulSourceFetch(ctx, sourceFetch, count); err != nil {
return fmt.Errorf("cannot commit successful source fetch: %w", err)
}
if err := w.finalizeCampaignFetchLifecycle(ctx, sourceFetch.TenantID, sourceFetch.AccessReviewCampaignID); err != nil {
if err := h.finalizeCampaignFetchLifecycle(ctx, sourceFetch.TenantID, sourceFetch.AccessReviewCampaignID); err != nil {
return fmt.Errorf("cannot finalize campaign fetch lifecycle: %w", err)
}
return nil
}
func (w *SourceFetchWorker) recoverStaleRows(ctx context.Context) {
now := time.Now()
staleThreshold := now.Add(-w.staleAfter)
err := w.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
var fetches coredata.AccessReviewCampaignSourceFetches
count, err := fetches.RecoverStale(ctx, tx, staleThreshold, now)
if err != nil {
return fmt.Errorf("cannot recover stale source fetches: %w", err)
}
if count > 0 {
w.logger.InfoCtx(
ctx,
"recovered stale source fetches",
log.Int64("count", count),
)
}
return nil
},
)
if err != nil {
w.logger.ErrorCtx(ctx, "cannot recover stale rows", log.Error(err))
}
}
func (w *SourceFetchWorker) commitFailedSourceFetch(
func (h *sourceFetchHandler) commitFailedSourceFetch(
ctx context.Context,
sourceFetch *coredata.AccessReviewCampaignSourceFetch,
failureErr error,
@@ -251,7 +178,7 @@ func (w *SourceFetchWorker) commitFailedSourceFetch(
sourceFetch.CompletedAt = new(now)
sourceFetch.UpdatedAt = now
return w.pg.WithTx(
return h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
return sourceFetch.Update(ctx, tx, scope)
@@ -259,7 +186,7 @@ func (w *SourceFetchWorker) commitFailedSourceFetch(
)
}
func (w *SourceFetchWorker) commitSuccessfulSourceFetch(
func (h *sourceFetchHandler) commitSuccessfulSourceFetch(
ctx context.Context,
sourceFetch *coredata.AccessReviewCampaignSourceFetch,
fetchedAccountsCount int,
@@ -275,7 +202,7 @@ func (w *SourceFetchWorker) commitSuccessfulSourceFetch(
sourceFetch.CompletedAt = new(now)
sourceFetch.UpdatedAt = now
return w.pg.WithTx(
return h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
return sourceFetch.Update(ctx, tx, scope)
@@ -283,14 +210,14 @@ func (w *SourceFetchWorker) commitSuccessfulSourceFetch(
)
}
func (w *SourceFetchWorker) finalizeCampaignFetchLifecycle(
func (h *sourceFetchHandler) finalizeCampaignFetchLifecycle(
ctx context.Context,
tenantID gid.TenantID,
campaignID gid.GID,
) error {
scope := coredata.NewScope(tenantID)
return w.pg.WithTx(
return h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
if err := lockCampaignForUpdate(ctx, tx, scope, campaignID); err != nil {

View File

@@ -18,12 +18,12 @@ import (
"context"
"errors"
"fmt"
"sync"
"time"
"go.gearno.de/crypto/uuid"
"go.gearno.de/kit/log"
"go.gearno.de/kit/pg"
"go.gearno.de/kit/worker"
"go.gearno.de/x/ref"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/filemanager"
@@ -37,42 +37,20 @@ import (
// organization that owns the given trust center.
type EmailPresenterConfigFunc func(ctx context.Context, organizationID gid.GID) (emails.PresenterConfig, error)
type (
CompletionCertificateWorker struct {
pg *pg.Client
fileManager *filemanager.Service
certificateGen *CertificateGenerator
presenterConfigFunc EmailPresenterConfigFunc
bucket string
logger *log.Logger
interval time.Duration
staleAfter time.Duration
maxConcurrency int
}
CompletionCertificateWorkerOption func(*CompletionCertificateWorker)
)
type completionCertificateHandler struct {
pg *pg.Client
fileManager *filemanager.Service
certificateGen *CertificateGenerator
presenterConfigFunc EmailPresenterConfigFunc
bucket string
logger *log.Logger
staleAfter time.Duration
}
const (
certificateFilename = "certificate-of-completion.pdf"
)
func WithCompletionCertificateWorkerInterval(d time.Duration) CompletionCertificateWorkerOption {
return func(w *CompletionCertificateWorker) { w.interval = d }
}
func WithCompletionCertificateWorkerStaleAfter(d time.Duration) CompletionCertificateWorkerOption {
return func(w *CompletionCertificateWorker) { w.staleAfter = d }
}
func WithCompletionCertificateWorkerMaxConcurrency(n int) CompletionCertificateWorkerOption {
return func(w *CompletionCertificateWorker) {
if n > 0 {
w.maxConcurrency = n
}
}
}
func NewCompletionCertificateWorker(
pgClient *pg.Client,
fileManager *filemanager.Service,
@@ -80,123 +58,93 @@ func NewCompletionCertificateWorker(
presenterConfigFunc EmailPresenterConfigFunc,
bucket string,
logger *log.Logger,
opts ...CompletionCertificateWorkerOption,
) *CompletionCertificateWorker {
w := &CompletionCertificateWorker{
opts ...worker.Option,
) *worker.Worker[coredata.ElectronicSignature] {
h := &completionCertificateHandler{
pg: pgClient,
fileManager: fileManager,
certificateGen: certificateGen,
presenterConfigFunc: presenterConfigFunc,
bucket: bucket,
logger: logger,
interval: 10 * time.Second,
staleAfter: 10 * time.Minute,
maxConcurrency: 5,
}
for _, opt := range opts {
opt(w)
}
return w
return worker.New(
"completion-certificate-worker",
h,
logger,
opts...,
)
}
func (w *CompletionCertificateWorker) Run(ctx context.Context) error {
var (
wg sync.WaitGroup
sem = make(chan struct{}, w.maxConcurrency)
)
defer wg.Wait()
func (h *completionCertificateHandler) Claim(ctx context.Context) (coredata.ElectronicSignature, error) {
var signature coredata.ElectronicSignature
LOOP:
select {
case <-ctx.Done():
return ctx.Err()
case <-time.After(w.interval):
// From there we should not accept cancelations anymore.
nonCancelableCtx := context.WithoutCancel(ctx)
w.recoverStaleCertificateRows(nonCancelableCtx)
for {
if err := w.processNext(ctx, sem, &wg); err != nil {
if !errors.Is(err, coredata.ErrResourceNotFound) {
w.logger.ErrorCtx(ctx, "cannot process certificate", log.Error(err))
}
break
}
}
goto LOOP
}
}
func (w *CompletionCertificateWorker) processNext(ctx context.Context, sem chan struct{}, wg *sync.WaitGroup) error {
select {
case sem <- struct{}{}:
case <-ctx.Done(): // FIXME: this will never be fired
return ctx.Err()
}
var (
signature coredata.ElectronicSignature
now = time.Now()
// From there we should not accept cancelations anymore.
nonCancelableCtx = context.WithoutCancel(ctx)
)
if err := w.pg.WithTx(
nonCancelableCtx,
if err := h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
if err := signature.LoadNextCompletedWithoutCertificateForUpdate(nonCancelableCtx, tx); err != nil {
if err := signature.LoadNextCompletedWithoutCertificateForUpdate(ctx, tx); err != nil {
return err
}
now := time.Now()
scope := coredata.NewScopeFromObjectID(signature.ID)
signature.CertificateProcessingStartedAt = &now
signature.AttemptCount++
signature.LastAttemptedAt = &now
signature.UpdatedAt = now
if err := signature.Update(nonCancelableCtx, tx, scope); err != nil {
if err := signature.Update(ctx, tx, scope); err != nil {
return fmt.Errorf("cannot update signature: %w", err)
}
return nil
},
); err != nil {
<-sem
return err
if errors.Is(err, coredata.ErrResourceNotFound) {
return coredata.ElectronicSignature{}, worker.ErrNoTask
}
return coredata.ElectronicSignature{}, err
}
wg.Add(1)
go func(signature coredata.ElectronicSignature) {
defer wg.Done()
defer func() { <-sem }()
return signature, nil
}
scope := coredata.NewScopeFromObjectID(signature.ID)
func (h *completionCertificateHandler) Process(ctx context.Context, signature coredata.ElectronicSignature) error {
scope := coredata.NewScopeFromObjectID(signature.ID)
if err := w.generateAndCommit(nonCancelableCtx, &signature); err != nil {
if err := w.handleCertFailure(nonCancelableCtx, &signature, scope, err); err != nil {
w.logger.ErrorCtx(nonCancelableCtx, "cannot handle certificate failure", log.Error(err))
}
if err := h.generateAndCommit(ctx, &signature); err != nil {
if err := h.handleCertFailure(ctx, &signature, scope, err); err != nil {
h.logger.ErrorCtx(ctx, "cannot handle certificate failure", log.Error(err))
}
}(signature)
return err
}
return nil
}
func (w *CompletionCertificateWorker) generateAndCommit(
func (h *completionCertificateHandler) RecoverStale(ctx context.Context) error {
return h.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) error {
return coredata.ResetStaleCertificateProcessing(ctx, conn, h.staleAfter)
},
)
}
func (h *completionCertificateHandler) generateAndCommit(
ctx context.Context,
signature *coredata.ElectronicSignature,
) error {
var (
scope = coredata.NewScopeFromObjectID(signature.ID)
)
scope := coredata.NewScopeFromObjectID(signature.ID)
email, attachments, err := w.generateCertificate(ctx, signature, scope)
email, attachments, err := h.generateCertificate(ctx, signature, scope)
if err != nil {
return err
}
if err := w.pg.WithTx(
if err := h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
signature.CertificateFileID = &attachments[1].FileID
@@ -232,7 +180,7 @@ func (w *CompletionCertificateWorker) generateAndCommit(
return nil
}
func (w *CompletionCertificateWorker) generateCertificate(
func (h *completionCertificateHandler) generateCertificate(
ctx context.Context,
signature *coredata.ElectronicSignature,
scope coredata.Scoper,
@@ -243,7 +191,7 @@ func (w *CompletionCertificateWorker) generateCertificate(
organization = coredata.Organization{}
)
if err := w.pg.WithConn(
if err := h.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) error {
if err := events.LoadBySignatureID(ctx, conn, scope, signature.ID); err != nil {
@@ -264,7 +212,7 @@ func (w *CompletionCertificateWorker) generateCertificate(
return nil, nil, err
}
certificatePDFReader, err := w.certificateGen.Generate(ctx, signature, events)
certificatePDFReader, err := h.certificateGen.Generate(ctx, signature, events)
if err != nil {
return nil, nil, fmt.Errorf("cannot generate certificate: %w", err)
}
@@ -272,7 +220,7 @@ func (w *CompletionCertificateWorker) generateCertificate(
certificateOfCompletionFile := coredata.File{
ID: gid.New(scope.GetTenantID(), coredata.FileEntityType),
OrganizationID: signature.OrganizationID,
BucketName: w.bucket,
BucketName: h.bucket,
MimeType: "application/pdf",
FileName: certificateFilename,
FileKey: uuid.MustNewV4().String(),
@@ -281,7 +229,7 @@ func (w *CompletionCertificateWorker) generateCertificate(
UpdatedAt: time.Now(),
}
certificateOfCompletionFileSize, err := w.fileManager.PutFile(
certificateOfCompletionFileSize, err := h.fileManager.PutFile(
ctx,
&certificateOfCompletionFile,
certificatePDFReader,
@@ -296,7 +244,7 @@ func (w *CompletionCertificateWorker) generateCertificate(
certificateOfCompletionFile.FileSize = certificateOfCompletionFileSize
if err := w.pg.WithTx(
if err := h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
if err := certificateOfCompletionFile.Insert(ctx, tx, scope); err != nil {
@@ -309,11 +257,11 @@ func (w *CompletionCertificateWorker) generateCertificate(
return nil, nil, err
}
presenterCfg, err := w.presenterConfigFunc(ctx, signature.OrganizationID)
presenterCfg, err := h.presenterConfigFunc(ctx, signature.OrganizationID)
if err != nil {
return nil, nil, fmt.Errorf("cannot resolve presenter config: %w", err)
}
emailPresenter := emails.NewPresenterFromConfig(w.fileManager, presenterCfg, ref.UnrefOrZero(signature.SignerFullName))
emailPresenter := emails.NewPresenterFromConfig(h.fileManager, presenterCfg, ref.UnrefOrZero(signature.SignerFullName))
docName := ref.UnrefOrZero(signature.DocumentName)
if docName == "" {
@@ -351,20 +299,20 @@ func (w *CompletionCertificateWorker) generateCertificate(
return email, attachments, nil
}
func (w *CompletionCertificateWorker) handleCertFailure(
func (h *completionCertificateHandler) handleCertFailure(
ctx context.Context,
signature *coredata.ElectronicSignature,
scope coredata.Scoper,
processingError error,
) error {
w.logger.ErrorCtx(
h.logger.ErrorCtx(
ctx,
"certificate worker failure",
log.Error(processingError),
log.String("signature_id", signature.ID.String()),
)
return w.pg.WithTx(
return h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
errStr := processingError.Error()
@@ -384,14 +332,3 @@ func (w *CompletionCertificateWorker) handleCertFailure(
},
)
}
func (w *CompletionCertificateWorker) recoverStaleCertificateRows(ctx context.Context) {
if err := w.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) error {
return coredata.ResetStaleCertificateProcessing(ctx, conn, w.staleAfter)
},
); err != nil {
w.logger.ErrorCtx(ctx, "cannot recover stale certificates", log.Error(err))
}
}

View File

@@ -18,11 +18,11 @@ import (
"context"
"errors"
"fmt"
"sync"
"time"
"go.gearno.de/kit/log"
"go.gearno.de/kit/pg"
"go.gearno.de/kit/worker"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/crypto/hash"
"go.probo.inc/probo/pkg/filemanager"
@@ -40,38 +40,24 @@ const (
)
type (
SealingWorker struct {
pg *pg.Client
fileManager *filemanager.Service
tsaClient *TSAClient
logger *log.Logger
interval time.Duration
tsaTimeout time.Duration
staleAfter time.Duration
maxConcurrency int
sealingHandler struct {
pg *pg.Client
fileManager *filemanager.Service
tsaClient *TSAClient
logger *log.Logger
tsaTimeout time.Duration
staleAfter time.Duration
}
SealingWorkerOption func(*SealingWorker)
SealingWorkerOption func(*sealingHandler)
)
func WithSealingWorkerInterval(d time.Duration) SealingWorkerOption {
return func(w *SealingWorker) { w.interval = d }
}
func WithSealingWorkerTSATimeout(d time.Duration) SealingWorkerOption {
return func(w *SealingWorker) { w.tsaTimeout = d }
return func(h *sealingHandler) { h.tsaTimeout = d }
}
func WithSealingWorkerStaleAfter(d time.Duration) SealingWorkerOption {
return func(w *SealingWorker) { w.staleAfter = d }
}
func WithSealingWorkerMaxConcurrency(n int) SealingWorkerOption {
return func(w *SealingWorker) {
if n > 0 {
w.maxConcurrency = n
}
}
return func(h *sealingHandler) { h.staleAfter = d }
}
func NewSealingWorker(
@@ -79,110 +65,82 @@ func NewSealingWorker(
fileManager *filemanager.Service,
tsaClient *TSAClient,
logger *log.Logger,
opts ...SealingWorkerOption,
) *SealingWorker {
w := &SealingWorker{
pg: pgClient,
fileManager: fileManager,
tsaClient: tsaClient,
logger: logger,
interval: 10 * time.Second,
tsaTimeout: 10 * time.Second,
staleAfter: 5 * time.Minute,
maxConcurrency: 5,
handlerOpts []SealingWorkerOption,
workerOpts ...worker.Option,
) *worker.Worker[coredata.ElectronicSignature] {
h := &sealingHandler{
pg: pgClient,
fileManager: fileManager,
tsaClient: tsaClient,
logger: logger,
tsaTimeout: 10 * time.Second,
staleAfter: 5 * time.Minute,
}
for _, opt := range opts {
opt(w)
for _, opt := range handlerOpts {
opt(h)
}
return w
return worker.New(
"sealing-worker",
h,
logger,
workerOpts...,
)
}
func (w *SealingWorker) Run(ctx context.Context) error {
var (
wg sync.WaitGroup
sem = make(chan struct{}, w.maxConcurrency)
)
func (h *sealingHandler) Claim(ctx context.Context) (coredata.ElectronicSignature, error) {
var signature coredata.ElectronicSignature
defer wg.Wait()
LOOP:
select {
case <-ctx.Done():
return ctx.Err()
case <-time.After(w.interval):
// From there we should not accept cancelations anymore.
nonCancelableCtx := context.WithoutCancel(ctx)
w.recoverStaleRows(nonCancelableCtx)
for {
if err := w.processNext(ctx, sem, &wg); err != nil {
if !errors.Is(err, coredata.ErrResourceNotFound) {
w.logger.ErrorCtx(nonCancelableCtx, "cannot claim signature", log.Error(err))
}
break
}
}
goto LOOP
}
}
func (w *SealingWorker) processNext(ctx context.Context, sem chan struct{}, wg *sync.WaitGroup) error {
select {
case sem <- struct{}{}:
case <-ctx.Done():
return ctx.Err()
}
var (
signature = coredata.ElectronicSignature{}
now = time.Now()
// From there we should not accept cancelations anymore.
nonCancelableCtx = context.WithoutCancel(ctx)
)
if err := w.pg.WithTx(
nonCancelableCtx,
if err := h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
if err := signature.LoadNextAcceptedForUpdateSkipLocked(nonCancelableCtx, tx); err != nil {
if err := signature.LoadNextAcceptedForUpdateSkipLocked(ctx, tx); err != nil {
return err
}
now := time.Now()
signature.Status = coredata.ElectronicSignatureStatusProcessing
signature.ProcessingStartedAt = &now
signature.AttemptCount++
signature.LastAttemptedAt = &now
signature.UpdatedAt = now
if err := signature.Update(nonCancelableCtx, tx, coredata.NewNoScope()); err != nil {
if err := signature.Update(ctx, tx, coredata.NewNoScope()); err != nil {
return fmt.Errorf("cannot update signature: %w", err)
}
return nil
},
); err != nil {
<-sem
return err
if errors.Is(err, coredata.ErrResourceNotFound) {
return coredata.ElectronicSignature{}, worker.ErrNoTask
}
return coredata.ElectronicSignature{}, err
}
wg.Add(1)
go func(signature coredata.ElectronicSignature) {
defer wg.Done()
defer func() { <-sem }()
return signature, nil
}
if err := w.sealAndCommit(nonCancelableCtx, &signature); err != nil {
if err := w.failSignature(nonCancelableCtx, &signature, err); err != nil {
w.logger.ErrorCtx(nonCancelableCtx, "cannot fail signature", log.Error(err))
}
func (h *sealingHandler) Process(ctx context.Context, signature coredata.ElectronicSignature) error {
if err := h.sealAndCommit(ctx, &signature); err != nil {
if err := h.failSignature(ctx, &signature, err); err != nil {
h.logger.ErrorCtx(ctx, "cannot fail signature", log.Error(err))
}
}(signature)
return err
}
return nil
}
func (w *SealingWorker) sealAndCommit(
func (h *sealingHandler) RecoverStale(ctx context.Context) error {
return h.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) error {
return coredata.ResetStaleProcessingSignatures(ctx, conn, h.staleAfter)
},
)
}
func (h *sealingHandler) sealAndCommit(
ctx context.Context,
signature *coredata.ElectronicSignature,
) error {
@@ -192,7 +150,7 @@ func (w *SealingWorker) sealAndCommit(
events []coredata.ElectronicSignatureEvent
)
if err := w.pg.WithConn(
if err := h.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) error {
if err := file.LoadByID(ctx, conn, scope, signature.FileID); err != nil {
@@ -205,7 +163,7 @@ func (w *SealingWorker) sealAndCommit(
return fmt.Errorf("%w: %w", ErrLoadFile, err)
}
pdfBytes, err := w.fileManager.GetFileBytes(ctx, &file)
pdfBytes, err := h.fileManager.GetFileBytes(ctx, &file)
if err != nil {
return fmt.Errorf("%w: %w", ErrDownloadPDF, err)
}
@@ -227,9 +185,9 @@ func (w *SealingWorker) sealAndCommit(
),
)
tsaCtx, cancel := context.WithTimeout(ctx, w.tsaTimeout)
tsaCtx, cancel := context.WithTimeout(ctx, h.tsaTimeout)
defer cancel()
tsaToken, err := w.tsaClient.Timestamp(tsaCtx, []byte(seal))
tsaToken, err := h.tsaClient.Timestamp(tsaCtx, []byte(seal))
if err != nil {
return fmt.Errorf("%w: %w", ErrTSATimestamp, err)
}
@@ -242,7 +200,7 @@ func (w *SealingWorker) sealAndCommit(
),
)
if err := w.pg.WithTx(
if err := h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
var current coredata.ElectronicSignature
@@ -282,21 +240,21 @@ func (w *SealingWorker) sealAndCommit(
return nil
}
func (w *SealingWorker) failSignature(
func (h *sealingHandler) failSignature(
ctx context.Context,
signature *coredata.ElectronicSignature,
processingError error,
) error {
scope := coredata.NewScopeFromObjectID(signature.ID)
w.logger.ErrorCtx(
h.logger.ErrorCtx(
ctx,
"sealing worker failure",
log.Error(processingError),
log.String("signature_id", signature.ID.String()),
)
return w.pg.WithTx(
return h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
errStr := userFacingError(processingError)
@@ -336,14 +294,3 @@ func userFacingError(err error) string {
return "An unexpected error occurred while processing your signature."
}
}
func (w *SealingWorker) recoverStaleRows(ctx context.Context) {
if err := w.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) error {
return coredata.ResetStaleProcessingSignatures(ctx, conn, w.staleAfter)
},
); err != nil {
w.logger.ErrorCtx(ctx, "cannot recover stale signatures", log.Error(err))
}
}

View File

@@ -111,18 +111,17 @@ func NewService(
func (s *Service) Run(ctx context.Context, presenterConfigFunc EmailPresenterConfigFunc) error {
g, gctx := errgroup.WithContext(ctx)
nonCancelableCtx := context.WithoutCancel(ctx)
sealingWorkerCtx, stopSealingWorker := context.WithCancel(nonCancelableCtx)
sealingWorkerCtx, stopSealingWorker := context.WithCancel(ctx)
sealingWorker := NewSealingWorker(
s.pg,
s.fileManager,
s.tsaClient,
s.logger.Named("sealing-worker"),
nil,
)
g.Go(func() error { return sealingWorker.Run(sealingWorkerCtx) })
certWorkerCtx, stopCertWorker := context.WithCancel(nonCancelableCtx)
certWorkerCtx, stopCertWorker := context.WithCancel(ctx)
certWorker := NewCompletionCertificateWorker(
s.pg,
s.fileManager,

View File

@@ -22,28 +22,26 @@ import (
"fmt"
"net"
"net/smtp"
"sync"
"time"
"github.com/jhillyerd/enmime"
"go.gearno.de/kit/log"
"go.gearno.de/kit/pg"
"go.gearno.de/kit/worker"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/filemanager"
)
type (
SendingWorker struct {
pg *pg.Client
fileManager *filemanager.Service
logger *log.Logger
smtp SMTPConfig
senderName string
senderEmail string
interval time.Duration
smtpTimeout time.Duration
staleAfter time.Duration
maxConcurrency int
sendingHandler struct {
pg *pg.Client
fileManager *filemanager.Service
logger *log.Logger
smtp SMTPConfig
senderName string
senderEmail string
smtpTimeout time.Duration
staleAfter time.Duration
}
SMTPConfig struct {
@@ -53,27 +51,15 @@ type (
TLSRequired bool
}
SendingWorkerOption func(*SendingWorker)
SendingWorkerOption func(*sendingHandler)
)
func WithSendingWorkerInterval(d time.Duration) SendingWorkerOption {
return func(w *SendingWorker) { w.interval = d }
}
func WithSendingWorkerSMTPTimeout(d time.Duration) SendingWorkerOption {
return func(w *SendingWorker) { w.smtpTimeout = d }
return func(h *sendingHandler) { h.smtpTimeout = d }
}
func WithSendingWorkerStaleAfter(d time.Duration) SendingWorkerOption {
return func(w *SendingWorker) { w.staleAfter = d }
}
func WithSendingWorkerMaxConcurrency(n int) SendingWorkerOption {
return func(w *SendingWorker) {
if n > 0 {
w.maxConcurrency = n
}
}
return func(h *sendingHandler) { h.staleAfter = d }
}
func NewSendingWorker(
@@ -83,75 +69,39 @@ func NewSendingWorker(
senderEmail string,
smtpCfg SMTPConfig,
logger *log.Logger,
opts ...SendingWorkerOption,
) *SendingWorker {
w := &SendingWorker{
pg: pgClient,
fileManager: fileManager,
logger: logger,
smtp: smtpCfg,
senderName: senderName,
senderEmail: senderEmail,
interval: 30 * time.Second,
smtpTimeout: 25 * time.Second,
staleAfter: 5 * time.Minute,
maxConcurrency: 20,
handlerOpts []SendingWorkerOption,
workerOpts ...worker.Option,
) *worker.Worker[coredata.Email] {
h := &sendingHandler{
pg: pgClient,
fileManager: fileManager,
logger: logger,
smtp: smtpCfg,
senderName: senderName,
senderEmail: senderEmail,
smtpTimeout: 25 * time.Second,
staleAfter: 5 * time.Minute,
}
for _, opt := range opts {
opt(w)
for _, opt := range handlerOpts {
opt(h)
}
return w
return worker.New(
"sending-worker",
h,
logger,
workerOpts...,
)
}
func (w *SendingWorker) Run(ctx context.Context) error {
var (
wg sync.WaitGroup
sem = make(chan struct{}, w.maxConcurrency)
)
func (h *sendingHandler) Claim(ctx context.Context) (coredata.Email, error) {
var email coredata.Email
defer wg.Wait()
ticker := time.NewTicker(w.interval)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return ctx.Err()
case <-ticker.C:
nonCancelableCtx := context.WithoutCancel(ctx)
w.recoverStaleRows(nonCancelableCtx)
for {
if err := w.processNext(ctx, sem, &wg); err != nil {
if !errors.Is(err, coredata.ErrNoUnsentEmail) {
w.logger.ErrorCtx(nonCancelableCtx, "cannot process email", log.Error(err))
}
break
}
}
}
}
}
func (w *SendingWorker) processNext(ctx context.Context, sem chan struct{}, wg *sync.WaitGroup) error {
select {
case sem <- struct{}{}:
case <-ctx.Done():
return ctx.Err()
}
var (
email = coredata.Email{}
nonCancelableCtx = context.WithoutCancel(ctx)
)
if err := w.pg.WithTx(
nonCancelableCtx,
if err := h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
if err := email.LoadNextPendingForUpdateSkipLocked(nonCancelableCtx, tx); err != nil {
if err := email.LoadNextPendingForUpdateSkipLocked(ctx, tx); err != nil {
return err
}
@@ -162,39 +112,48 @@ func (w *SendingWorker) processNext(ctx context.Context, sem chan struct{}, wg *
email.LastAttemptedAt = &now
email.UpdatedAt = now
if err := email.Update(nonCancelableCtx, tx); err != nil {
if err := email.Update(ctx, tx); err != nil {
return fmt.Errorf("cannot update email: %w", err)
}
return nil
},
); err != nil {
<-sem
return err
if errors.Is(err, coredata.ErrNoUnsentEmail) {
return coredata.Email{}, worker.ErrNoTask
}
return coredata.Email{}, err
}
wg.Add(1)
go func(email coredata.Email) {
defer wg.Done()
defer func() { <-sem }()
return email, nil
}
if sendErr := w.sendAndCommit(nonCancelableCtx, &email); sendErr != nil {
if failErr := w.failEmail(nonCancelableCtx, &email, sendErr); failErr != nil {
w.logger.ErrorCtx(nonCancelableCtx, "cannot fail email", log.Error(failErr))
}
func (h *sendingHandler) Process(ctx context.Context, email coredata.Email) error {
if sendErr := h.sendAndCommit(ctx, &email); sendErr != nil {
if failErr := h.failEmail(ctx, &email, sendErr); failErr != nil {
h.logger.ErrorCtx(ctx, "cannot fail email", log.Error(failErr))
}
}(email)
return sendErr
}
return nil
}
func (w *SendingWorker) sendAndCommit(
func (h *sendingHandler) RecoverStale(ctx context.Context) error {
return h.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) error {
return coredata.ResetStaleProcessingEmails(ctx, conn, h.staleAfter)
},
)
}
func (h *sendingHandler) sendAndCommit(
ctx context.Context,
email *coredata.Email,
) error {
var buf bytes.Buffer
if err := w.pg.WithConn(
if err := h.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) error {
var attachments coredata.EmailAttachments
@@ -202,14 +161,14 @@ func (w *SendingWorker) sendAndCommit(
return fmt.Errorf("cannot load email attachments: %w", err)
}
fromName := w.senderName
fromName := h.senderName
if email.SenderName != nil {
fromName = *email.SenderName + " via " + w.senderName
fromName = *email.SenderName + " via " + h.senderName
}
mail := enmime.Builder().
Subject(email.Subject).
From(fromName, w.senderEmail).
From(fromName, h.senderEmail).
To(email.RecipientName, email.RecipientEmail).
Text([]byte(email.TextBody))
@@ -233,7 +192,7 @@ func (w *SendingWorker) sendAndCommit(
return fmt.Errorf("cannot load file record for attachment %s: %w", att.Filename, err)
}
data, err := w.fileManager.GetFileBytes(ctx, &file)
data, err := h.fileManager.GetFileBytes(ctx, &file)
if err != nil {
return fmt.Errorf("cannot download attachment %s: %w", att.Filename, err)
}
@@ -256,17 +215,17 @@ func (w *SendingWorker) sendAndCommit(
return err
}
sendCtx, cancel := context.WithTimeout(ctx, w.smtpTimeout)
sendCtx, cancel := context.WithTimeout(ctx, h.smtpTimeout)
defer cancel()
if err := w.sendMail(sendCtx, []string{email.RecipientEmail}, buf.Bytes()); err != nil {
if err := h.sendMail(sendCtx, []string{email.RecipientEmail}, buf.Bytes()); err != nil {
if errors.Is(err, context.DeadlineExceeded) {
return fmt.Errorf("email sending timed out after %s: %w", w.smtpTimeout, err)
return fmt.Errorf("email sending timed out after %s: %w", h.smtpTimeout, err)
}
return fmt.Errorf("cannot send email: %w", err)
}
if err := w.pg.WithTx(
if err := h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
now := time.Now()
@@ -283,7 +242,7 @@ func (w *SendingWorker) sendAndCommit(
return nil
},
); err != nil {
w.logger.ErrorCtx(ctx,
h.logger.ErrorCtx(ctx,
"email sent but failed to commit status update; will not re-queue to avoid duplicate delivery",
log.Error(err),
log.String("email_id", email.ID.String()),
@@ -293,19 +252,19 @@ func (w *SendingWorker) sendAndCommit(
return nil
}
func (w *SendingWorker) failEmail(
func (h *sendingHandler) failEmail(
ctx context.Context,
email *coredata.Email,
processingError error,
) error {
w.logger.ErrorCtx(
h.logger.ErrorCtx(
ctx,
"sending worker failure",
log.Error(processingError),
log.String("email_id", email.ID.String()),
)
return w.pg.WithTx(
return h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
errStr := processingError.Error()
@@ -328,26 +287,15 @@ func (w *SendingWorker) failEmail(
)
}
func (w *SendingWorker) recoverStaleRows(ctx context.Context) {
if err := w.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) error {
return coredata.ResetStaleProcessingEmails(ctx, conn, w.staleAfter)
},
); err != nil {
w.logger.ErrorCtx(ctx, "cannot recover stale emails", log.Error(err))
}
}
func (w *SendingWorker) sendMail(ctx context.Context, to []string, msg []byte) error {
host, _, err := net.SplitHostPort(w.smtp.Addr)
func (h *sendingHandler) sendMail(ctx context.Context, to []string, msg []byte) error {
host, _, err := net.SplitHostPort(h.smtp.Addr)
if err != nil {
return fmt.Errorf("invalid address: %w", err)
}
var d net.Dialer
conn, err := d.DialContext(ctx, "tcp", w.smtp.Addr)
conn, err := d.DialContext(ctx, "tcp", h.smtp.Addr)
if err != nil {
return fmt.Errorf("connection error: %w", err)
}
@@ -365,20 +313,20 @@ func (w *SendingWorker) sendMail(ctx context.Context, to []string, msg []byte) e
}
defer func() { _ = c.Quit() }()
if w.smtp.TLSRequired {
if h.smtp.TLSRequired {
if err := c.StartTLS(&tls.Config{ServerName: host}); err != nil {
return fmt.Errorf("TLS negotiation error: %w", err)
}
}
if w.smtp.User != "" && w.smtp.Password != "" {
auth := smtp.PlainAuth("", w.smtp.User, w.smtp.Password, host)
if h.smtp.User != "" && h.smtp.Password != "" {
auth := smtp.PlainAuth("", h.smtp.User, h.smtp.Password, host)
if err = c.Auth(auth); err != nil {
return fmt.Errorf("SMTP authentication error: %w", err)
}
}
if err = c.Mail(w.senderEmail); err != nil {
if err = c.Mail(h.senderEmail); err != nil {
return fmt.Errorf("MAIL FROM error: %w", err)
}

View File

@@ -18,163 +18,114 @@ import (
"context"
"errors"
"fmt"
"sync"
"time"
"go.gearno.de/kit/log"
"go.gearno.de/kit/pg"
"go.gearno.de/kit/worker"
"go.probo.inc/probo/pkg/coredata"
)
type (
MailingListWorker struct {
service *Service
pg *pg.Client
logger *log.Logger
interval time.Duration
staleAfter time.Duration
maxConcurrency int
}
MailingListWorkerOption func(*MailingListWorker)
)
func WithMailingListWorkerInterval(d time.Duration) MailingListWorkerOption {
return func(w *MailingListWorker) { w.interval = d }
}
func WithMailingListWorkerStaleAfter(d time.Duration) MailingListWorkerOption {
return func(w *MailingListWorker) { w.staleAfter = d }
}
func WithMailingListWorkerMaxConcurrency(n int) MailingListWorkerOption {
return func(w *MailingListWorker) {
if n > 0 {
w.maxConcurrency = n
}
}
type mailingListHandler struct {
service *Service
pg *pg.Client
logger *log.Logger
staleAfter time.Duration
}
func NewMailingListWorker(
service *Service,
pgClient *pg.Client,
logger *log.Logger,
opts ...MailingListWorkerOption,
) *MailingListWorker {
w := &MailingListWorker{
service: service,
pg: pgClient,
logger: logger,
interval: 10 * time.Second,
staleAfter: 5 * time.Minute,
maxConcurrency: 5,
opts ...worker.Option,
) *worker.Worker[coredata.MailingListUpdate] {
h := &mailingListHandler{
service: service,
pg: pgClient,
logger: logger,
staleAfter: 5 * time.Minute,
}
for _, opt := range opts {
opt(w)
}
return w
return worker.New(
"mailing-list-worker",
h,
logger,
opts...,
)
}
func (w *MailingListWorker) Run(ctx context.Context) error {
var (
wg sync.WaitGroup
sem = make(chan struct{}, w.maxConcurrency)
)
func (h *mailingListHandler) Claim(ctx context.Context) (coredata.MailingListUpdate, error) {
var mlu coredata.MailingListUpdate
defer wg.Wait()
LOOP:
select {
case <-ctx.Done():
return ctx.Err()
case <-time.After(w.interval):
// From there we should not accept cancellations anymore.
nonCancelableCtx := context.WithoutCancel(ctx)
w.recoverStaleRows(nonCancelableCtx)
for {
if err := w.processNext(ctx, sem, &wg); err != nil {
if !errors.Is(err, coredata.ErrResourceNotFound) {
w.logger.ErrorCtx(nonCancelableCtx, "cannot claim mailing list update", log.Error(err))
}
break
}
}
goto LOOP
}
}
func (w *MailingListWorker) processNext(ctx context.Context, sem chan struct{}, wg *sync.WaitGroup) error {
select {
case sem <- struct{}{}:
case <-ctx.Done():
return ctx.Err()
}
var (
mlu coredata.MailingListUpdate
now = time.Now()
nonCancelableCtx = context.WithoutCancel(ctx)
)
if err := w.pg.WithTx(
nonCancelableCtx,
if err := h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
if err := mlu.LoadNextEnqueuedForUpdateSkipLocked(nonCancelableCtx, tx); err != nil {
if err := mlu.LoadNextEnqueuedForUpdateSkipLocked(ctx, tx); err != nil {
return err
}
scope := coredata.NewScopeFromObjectID(mlu.ID)
mlu.Status = coredata.MailingListUpdateStatusProcessing
mlu.UpdatedAt = now
mlu.UpdatedAt = time.Now()
if err := mlu.Update(nonCancelableCtx, tx, scope); err != nil {
if err := mlu.Update(ctx, tx, scope); err != nil {
return fmt.Errorf("cannot claim mailing list update: %w", err)
}
return nil
},
); err != nil {
<-sem
return err
if errors.Is(err, coredata.ErrResourceNotFound) {
return coredata.MailingListUpdate{}, worker.ErrNoTask
}
return coredata.MailingListUpdate{}, err
}
wg.Add(1)
go func(mlu coredata.MailingListUpdate) {
defer wg.Done()
defer func() { <-sem }()
return mlu, nil
}
if err := w.sendAndCommit(nonCancelableCtx, &mlu); err != nil {
w.logger.ErrorCtx(
nonCancelableCtx,
"cannot send mailing list update",
func (h *mailingListHandler) Process(ctx context.Context, mlu coredata.MailingListUpdate) error {
if err := h.sendAndCommit(ctx, &mlu); err != nil {
h.logger.ErrorCtx(
ctx,
"cannot send mailing list update",
log.Error(err),
log.String("mailing_list_update_id", mlu.ID.String()),
)
if err := h.resetEnqueued(ctx, &mlu); err != nil {
h.logger.ErrorCtx(
ctx,
"cannot reset mailing list update to enqueued",
log.Error(err),
log.String("mailing_list_update_id", mlu.ID.String()),
)
if err := w.resetEnqueued(nonCancelableCtx, &mlu); err != nil {
w.logger.ErrorCtx(
nonCancelableCtx,
"cannot reset mailing list update to enqueued",
log.Error(err),
log.String("mailing_list_update_id", mlu.ID.String()),
)
}
}
}(mlu)
return err
}
return nil
}
func (w *MailingListWorker) sendAndCommit(ctx context.Context, mlu *coredata.MailingListUpdate) error {
if err := w.service.CreateUpdateEmails(ctx, mlu.MailingListID, mlu.ID, mlu.Title, mlu.Body); err != nil {
func (h *mailingListHandler) RecoverStale(ctx context.Context) error {
return h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
if err := coredata.ResetStaleProcessingMailingListUpdates(ctx, tx, h.staleAfter); err != nil {
return fmt.Errorf("cannot reset stale processing mailing list updates: %w", err)
}
return nil
},
)
}
func (h *mailingListHandler) sendAndCommit(ctx context.Context, mlu *coredata.MailingListUpdate) error {
if err := h.service.CreateUpdateEmails(ctx, mlu.MailingListID, mlu.ID, mlu.Title, mlu.Body); err != nil {
return fmt.Errorf("cannot create update emails: %w", err)
}
return w.pg.WithTx(
return h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
scope := coredata.NewScopeFromObjectID(mlu.ID)
@@ -200,8 +151,8 @@ func (w *MailingListWorker) sendAndCommit(ctx context.Context, mlu *coredata.Mai
)
}
func (w *MailingListWorker) resetEnqueued(ctx context.Context, mlu *coredata.MailingListUpdate) error {
return w.pg.WithTx(
func (h *mailingListHandler) resetEnqueued(ctx context.Context, mlu *coredata.MailingListUpdate) error {
return h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
scope := coredata.NewScopeFromObjectID(mlu.ID)
@@ -216,19 +167,3 @@ func (w *MailingListWorker) resetEnqueued(ctx context.Context, mlu *coredata.Mai
},
)
}
func (w *MailingListWorker) recoverStaleRows(ctx context.Context) {
err := w.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
if err := coredata.ResetStaleProcessingMailingListUpdates(ctx, tx, w.staleAfter); err != nil {
return fmt.Errorf("cannot reset stale processing mailing list updates: %w", err)
}
return nil
},
)
if err != nil {
w.logger.ErrorCtx(ctx, "cannot recover stale processing mailing list updates", log.Error(err))
}
}

View File

@@ -18,159 +18,121 @@ import (
"context"
"errors"
"fmt"
"sync"
"time"
"go.gearno.de/kit/log"
"go.gearno.de/kit/pg"
"go.gearno.de/kit/worker"
"go.probo.inc/probo/pkg/coredata"
"go.probo.inc/probo/pkg/evidencedescriber"
"go.probo.inc/probo/pkg/filemanager"
)
type (
EvidenceDescriptionWorker struct {
pg *pg.Client
fileManager *filemanager.Service
describer *evidencedescriber.Describer
logger *log.Logger
interval time.Duration
staleAfter time.Duration
maxConcurrency int
evidenceDescriptionHandler struct {
pg *pg.Client
fileManager *filemanager.Service
describer *evidencedescriber.Describer
logger *log.Logger
staleAfter time.Duration
}
EvidenceDescriptionWorkerOption func(*EvidenceDescriptionWorker)
EvidenceDescriptionWorkerConfig struct {
StaleAfter time.Duration
}
)
func WithEvidenceDescriptionWorkerInterval(d time.Duration) EvidenceDescriptionWorkerOption {
return func(w *EvidenceDescriptionWorker) { w.interval = d }
}
func WithEvidenceDescriptionWorkerStaleAfter(d time.Duration) EvidenceDescriptionWorkerOption {
return func(w *EvidenceDescriptionWorker) { w.staleAfter = d }
}
func WithEvidenceDescriptionWorkerMaxConcurrency(n int) EvidenceDescriptionWorkerOption {
return func(w *EvidenceDescriptionWorker) {
if n > 0 {
w.maxConcurrency = n
}
}
}
func NewEvidenceDescriptionWorker(
pgClient *pg.Client,
fileManager *filemanager.Service,
describer *evidencedescriber.Describer,
logger *log.Logger,
opts ...EvidenceDescriptionWorkerOption,
) *EvidenceDescriptionWorker {
w := &EvidenceDescriptionWorker{
pg: pgClient,
fileManager: fileManager,
describer: describer,
logger: logger,
interval: 10 * time.Second,
staleAfter: 5 * time.Minute,
maxConcurrency: 10,
cfg EvidenceDescriptionWorkerConfig,
opts ...worker.Option,
) *worker.Worker[coredata.Evidence] {
staleAfter := cfg.StaleAfter
if staleAfter == 0 {
staleAfter = 5 * time.Minute
}
for _, opt := range opts {
opt(w)
h := &evidenceDescriptionHandler{
pg: pgClient,
fileManager: fileManager,
describer: describer,
logger: logger,
staleAfter: staleAfter,
}
return w
return worker.New(
"evidence-description-worker",
h,
logger,
opts...,
)
}
func (w *EvidenceDescriptionWorker) Run(ctx context.Context) error {
var (
wg sync.WaitGroup
sem = make(chan struct{}, w.maxConcurrency)
ticker = time.NewTicker(w.interval)
)
defer ticker.Stop()
defer wg.Wait()
func (h *evidenceDescriptionHandler) Claim(ctx context.Context) (coredata.Evidence, error) {
var evidence coredata.Evidence
for {
select {
case <-ctx.Done():
return ctx.Err()
case <-ticker.C:
nonCancelableCtx := context.WithoutCancel(ctx)
w.recoverStaleRows(nonCancelableCtx)
for {
if err := w.processNext(ctx, sem, &wg); err != nil {
if !errors.Is(err, coredata.ErrResourceNotFound) {
w.logger.ErrorCtx(nonCancelableCtx, "cannot claim evidence for description", log.Error(err))
}
break
}
}
}
}
}
func (w *EvidenceDescriptionWorker) processNext(ctx context.Context, sem chan struct{}, wg *sync.WaitGroup) error {
select {
case sem <- struct{}{}:
case <-ctx.Done():
return ctx.Err()
}
var (
evidence = coredata.Evidence{}
now = time.Now()
nonCancelableCtx = context.WithoutCancel(ctx)
)
if err := w.pg.WithTx(
nonCancelableCtx,
if err := h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
if err := evidence.LoadNextPendingDescriptionForUpdateSkipLocked(
nonCancelableCtx,
tx,
); err != nil {
if err := evidence.LoadNextPendingDescriptionForUpdateSkipLocked(ctx, tx); err != nil {
return err
}
now := time.Now()
evidence.DescriptionStatus = coredata.EvidenceDescriptionStatusProcessing
evidence.DescriptionProcessingStartedAt = &now
evidence.UpdatedAt = now
if err := evidence.Update(nonCancelableCtx, tx, coredata.NewNoScope()); err != nil {
if err := evidence.Update(ctx, tx, coredata.NewNoScope()); err != nil {
return fmt.Errorf("cannot update evidence: %w", err)
}
return nil
},
); err != nil {
<-sem
return err
if errors.Is(err, coredata.ErrResourceNotFound) {
return coredata.Evidence{}, worker.ErrNoTask
}
return coredata.Evidence{}, err
}
wg.Add(1)
go func(evidence coredata.Evidence) {
defer wg.Done()
defer func() { <-sem }()
return evidence, nil
}
if err := w.describeAndCommit(nonCancelableCtx, &evidence); err != nil {
w.logger.ErrorCtx(
nonCancelableCtx,
"evidence description worker failure",
log.Error(err),
log.String("evidence_id", evidence.ID.String()),
)
func (h *evidenceDescriptionHandler) Process(ctx context.Context, evidence coredata.Evidence) error {
if err := h.describeAndCommit(ctx, &evidence); err != nil {
h.logger.ErrorCtx(
ctx,
"evidence description worker failure",
log.Error(err),
log.String("evidence_id", evidence.ID.String()),
)
if err := w.failEvidence(nonCancelableCtx, &evidence); err != nil {
w.logger.ErrorCtx(nonCancelableCtx, "cannot mark evidence description as failed", log.Error(err))
}
if err := h.failEvidence(ctx, &evidence); err != nil {
h.logger.ErrorCtx(ctx, "cannot mark evidence description as failed", log.Error(err))
}
}(evidence)
return err
}
return nil
}
func (w *EvidenceDescriptionWorker) describeAndCommit(
func (h *evidenceDescriptionHandler) RecoverStale(ctx context.Context) error {
return h.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) error {
if err := coredata.ResetStaleDescriptionProcessing(ctx, conn, h.staleAfter); err != nil {
return fmt.Errorf("cannot reset stale description processing: %w", err)
}
return nil
},
)
}
func (h *evidenceDescriptionHandler) describeAndCommit(
ctx context.Context,
evidence *coredata.Evidence,
) error {
@@ -181,7 +143,7 @@ func (w *EvidenceDescriptionWorker) describeAndCommit(
scope := coredata.NewScopeFromObjectID(evidence.ID)
var file coredata.File
if err := w.pg.WithConn(
if err := h.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) error {
if err := file.LoadByID(ctx, conn, scope, *evidence.EvidenceFileId); err != nil {
@@ -193,17 +155,17 @@ func (w *EvidenceDescriptionWorker) describeAndCommit(
return fmt.Errorf("cannot load file: %w", err)
}
base64Data, mimeType, err := w.fileManager.GetFileBase64(ctx, &file)
base64Data, mimeType, err := h.fileManager.GetFileBase64(ctx, &file)
if err != nil {
return fmt.Errorf("cannot download file: %w", err)
}
description, err := w.describer.Describe(ctx, file.FileName, mimeType, base64Data)
description, err := h.describer.Describe(ctx, file.FileName, mimeType, base64Data)
if err != nil {
return fmt.Errorf("cannot describe evidence: %w", err)
}
return w.pg.WithTx(
return h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
evidence.Description = description
@@ -219,13 +181,13 @@ func (w *EvidenceDescriptionWorker) describeAndCommit(
)
}
func (w *EvidenceDescriptionWorker) failEvidence(
func (h *evidenceDescriptionHandler) failEvidence(
ctx context.Context,
evidence *coredata.Evidence,
) error {
scope := coredata.NewScopeFromObjectID(evidence.ID)
return w.pg.WithTx(
return h.pg.WithTx(
ctx,
func(ctx context.Context, tx pg.Tx) error {
evidence.DescriptionStatus = coredata.EvidenceDescriptionStatusFailed
@@ -239,17 +201,3 @@ func (w *EvidenceDescriptionWorker) failEvidence(
},
)
}
func (w *EvidenceDescriptionWorker) recoverStaleRows(ctx context.Context) {
if err := w.pg.WithConn(
ctx,
func(ctx context.Context, conn pg.Querier) error {
if err := coredata.ResetStaleDescriptionProcessing(ctx, conn, w.staleAfter); err != nil {
return fmt.Errorf("cannot reset stale description processing: %w", err)
}
return nil
},
); err != nil {
w.logger.ErrorCtx(ctx, "cannot recover stale evidence descriptions", log.Error(err))
}
}

View File

@@ -41,6 +41,7 @@ import (
"go.gearno.de/kit/migrator"
"go.gearno.de/kit/pg"
"go.gearno.de/kit/unit"
"go.gearno.de/kit/worker"
"go.opentelemetry.io/otel/trace"
"go.probo.inc/probo/pkg/accessreview"
"go.probo.inc/probo/pkg/awsconfig"
@@ -580,8 +581,11 @@ func (impl *Implm) Run(
TLSRequired: impl.cfg.Notifications.Mailer.SMTP.TLSRequired,
},
l.Named("sending-worker"),
mailer.WithSendingWorkerSMTPTimeout(time.Second*10),
mailer.WithSendingWorkerInterval(time.Duration(impl.cfg.Notifications.Mailer.MailerInterval)*time.Second),
[]mailer.SendingWorkerOption{
mailer.WithSendingWorkerSMTPTimeout(time.Second * 10),
},
worker.WithInterval(time.Duration(impl.cfg.Notifications.Mailer.MailerInterval)*time.Second),
worker.WithMaxConcurrency(20),
)
wg.Go(
func() {
@@ -676,9 +680,11 @@ func (impl *Implm) Run(
fileManagerService,
evidenceDescriber,
l.Named("evidence-description-worker"),
probo.WithEvidenceDescriptionWorkerInterval(time.Duration(impl.cfg.EvidenceDescriber.Interval)*time.Second),
probo.WithEvidenceDescriptionWorkerStaleAfter(time.Duration(impl.cfg.EvidenceDescriber.StaleAfter)*time.Second),
probo.WithEvidenceDescriptionWorkerMaxConcurrency(impl.cfg.EvidenceDescriber.MaxConcurrency),
probo.EvidenceDescriptionWorkerConfig{
StaleAfter: time.Duration(impl.cfg.EvidenceDescriber.StaleAfter) * time.Second,
},
worker.WithInterval(time.Duration(impl.cfg.EvidenceDescriber.Interval)*time.Second),
worker.WithMaxConcurrency(impl.cfg.EvidenceDescriber.MaxConcurrency),
)
evidenceDescriptionWorkerCtx, stopEvidenceDescriptionWorker := context.WithCancel(context.Background())
wg.Go(