feat: Implement storage job limiter

internal/pkg/service/stream/sink/type/tablesink/keboola/bridge/job.go

@@ -3,21 +3,20 @@ package bridge
 import (
 	"context"
 	"fmt"
+	"math"
+	"strconv"
+	"time"
 
 	"github.com/keboola/go-client/pkg/keboola"
 	"go.opentelemetry.io/otel/attribute"
 
 	"github.com/keboola/keboola-as-code/internal/pkg/service/common/ctxattr"
 	"github.com/keboola/keboola-as-code/internal/pkg/service/stream/definition/key"
 	"github.com/keboola/keboola-as-code/internal/pkg/service/stream/plugin"
-	keboolasink "github.com/keboola/keboola-as-code/internal/pkg/service/stream/sink/type/tablesink/keboola"
 	"github.com/keboola/keboola-as-code/internal/pkg/service/stream/sink/type/tablesink/keboola/bridge/model"
+	"github.com/keboola/keboola-as-code/internal/pkg/utils/errors"
 )
 
-func (b *Bridge) SinkToken(ctx context.Context, sinkKey key.SinkKey) (keboolasink.Token, error) {
-	return b.schema.Token().ForSink(sinkKey).GetOrErr(b.client).Do(ctx).ResultOrErr()
-}
-
 func (b *Bridge) createJob(ctx context.Context, file plugin.File, storageJob *keboola.StorageJob) error {
 	keboolaJob := model.Job{
 		JobKey: key.JobKey{SinkKey: file.SinkKey, JobID: key.JobID(storageJob.ID.String())},
@@ -58,3 +57,65 @@ func (b *Bridge) canAcceptNewFile(ctx context.Context, sinkKey key.SinkKey) bool
 
 	return runningJobs < b.config.JobLimit
 }
+
+func (b *Bridge) CleanJob(ctx context.Context, job model.Job) (err error, deleted bool) {
+	// Parse storage job ID from string
+	id64, err := strconv.ParseInt(string(job.JobKey.JobID), 10, 64)
+	if err != nil {
+		err = errors.PrefixErrorf(err, `cannot get keboola storage job "%s"`, job.JobKey)
+		b.logger.Error(ctx, err.Error())
+		return err, false
+	}
+
+	if id64 < math.MinInt || id64 > math.MaxInt {
+		err = errors.Errorf("parsed job ID %d is out of int range", id64)
+		b.logger.Error(ctx, err.Error())
+		return err, false
+	}
+
+	token, err := b.schema.Token().ForSink(job.SinkKey).GetOrErr(b.client).Do(ctx).ResultOrErr()
+	if err != nil {
+		b.logger.Warnf(ctx, "cannot get token for sink, already deleted: %s", err.Error())
+		return nil, false
+	}
+
+	// Get job details from storage API
+	id := int(id64)
+	api := b.publicAPI.NewAuthorizedAPI(token.TokenString(), 1*time.Minute)
+	var jobStatus *keboola.StorageJob
+	if jobStatus, err = api.GetStorageJobRequest(keboola.StorageJobKey{ID: keboola.StorageJobID(id)}).Send(ctx); err != nil {
+		b.logger.Warnf(ctx, "cannot get information about storage job, probably already deleted: %s", err.Error())
+		return nil, false
+	}
+
+	attr := attribute.String("job.state", jobStatus.Status)
+	ctx = ctxattr.ContextWith(ctx, attr)
+	// Check status of storage job
+	if jobStatus.Status == keboola.StorageJobStatusProcessing || jobStatus.Status == keboola.StorageJobStatusWaiting {
+		b.logger.Debugf(ctx, "cannot remove storage job, job status: %s", jobStatus.Status)
+		return nil, false
+	}
+
+	// Acquire lock
+	mutex := b.locks.NewMutex(fmt.Sprintf("api.source.sink.jobs.%s", job.SinkKey))
+	if err = mutex.TryLock(ctx); err != nil {
+		return err, false
+	}
+	defer func() {
+		if err := mutex.Unlock(ctx); err != nil {
+			b.logger.Errorf(ctx, "cannot unlock the lock: %s", err)
+		}
+	}()
+
+	// Purge job in bridge repository
+	if _, err = b.keboolaBridgeRepository.Job().Purge(&job).RequireLock(mutex).Do(ctx).ResultOrErr(); err != nil {
+		err = errors.PrefixErrorf(err, `cannot delete finished storage job "%s"`, job.JobKey)
+		b.logger.Error(ctx, err.Error())
+		return err, false
+	}
+
+	// Log job details
+	b.logger.Infof(ctx, `deleted finished storage job`)
+
+	return nil, true
+}

internal/pkg/service/stream/storage/metacleanup/metacleanup.go

@@ -4,9 +4,6 @@ package metacleanup
 
 import (
 	"context"
-	"fmt"
-	"math"
-	"strconv"
 	"sync"
 	"time"
 
@@ -176,10 +173,24 @@ func (n *Node) cleanMetadata(ctx context.Context) (err error) {
 		ListAll().
 		ForEach(func(job keboolaBridgeModel.Job, _ *iterator.Header) error {
 			grp.Go(func() error {
-				err, deleted := n.cleanJob(ctx, job)
+				// There can be several cleanup nodes, each node processes an own part.
+				if !n.dist.MustCheckIsOwner(job.ProjectID.String()) {
+					return nil
+				}
+
+				// Log/trace job details
+				attrs := job.Telemetry()
+				ctx = ctxattr.ContextWith(ctx, attrs...)
+
+				// Trace each job
+				ctx, span := n.telemetry.Tracer().Start(ctx, "keboola.go.stream.model.cleanup.metadata.cleanJob")
+				defer span.End(&err)
+
+				err, deleted := n.bridge.CleanJob(ctx, job)
 				if deleted {
 					jobCounter.Add(1)
 				}
+
 				return err
 			})
 
@@ -251,78 +262,3 @@ func (n *Node) isFileExpired(file model.File, age time.Duration) bool {
 	// Other files have a longer expiration so there is time for retries.
 	return age >= n.config.ActiveFileExpiration
 }
-
-func (n *Node) cleanJob(ctx context.Context, job keboolaBridgeModel.Job) (err error, deleted bool) {
-	// There can be several cleanup nodes, each node processes an own part.
-	if !n.dist.MustCheckIsOwner(job.ProjectID.String()) {
-		return nil, false
-	}
-
-	// Log/trace file details
-	attrs := job.Telemetry()
-	ctx = ctxattr.ContextWith(ctx, attrs...)
-
-	// Trace each job
-	ctx, span := n.telemetry.Tracer().Start(ctx, "keboola.go.stream.model.cleanup.metadata.cleanJob")
-	defer span.End(&err)
-
-	// Parse storage job ID from string
-	id64, err := strconv.ParseInt(string(job.JobKey.JobID), 10, 64)
-	if err != nil {
-		err = errors.PrefixErrorf(err, `cannot get keboola storage job "%s"`, job.JobKey)
-		n.logger.Error(ctx, err.Error())
-		return err, false
-	}
-
-	if id64 < math.MinInt || id64 > math.MaxInt {
-		err = errors.Errorf("parsed job ID %d is out of int range", id64)
-		n.logger.Error(ctx, err.Error())
-		return err, false
-	}
-
-	token, err := n.bridge.SinkToken(ctx, job.SinkKey)
-	if err != nil {
-		n.logger.Warnf(ctx, "cannot get token for sink, already deleted: %s", err.Error())
-		return nil, false
-	}
-
-	// Get job details from storage API
-	id := int(id64)
-	api := n.publicAPI.NewAuthorizedAPI(token.TokenString(), 1*time.Minute)
-	var jobStatus *keboola.StorageJob
-	if jobStatus, err = api.GetStorageJobRequest(keboola.StorageJobKey{ID: keboola.StorageJobID(id)}).Send(ctx); err != nil {
-		n.logger.Warnf(ctx, "cannot get information about storage job, probably already deleted: %s", err.Error())
-		return nil, false
-	}
-
-	attrs = append(attrs, attribute.String("job.state", jobStatus.Status))
-	ctx = ctxattr.ContextWith(ctx, attrs...)
-	// Check status of storage job
-	if jobStatus.Status == keboola.StorageJobStatusProcessing || jobStatus.Status == keboola.StorageJobStatusWaiting {
-		n.logger.Debugf(ctx, "cannot remove storage job, job status: %s", jobStatus.Status)
-		return nil, false
-	}
-
-	// Acquire lock
-	mutex := n.locks.NewMutex(fmt.Sprintf("api.source.sink.jobs.%s", job.SinkKey))
-	if err = mutex.TryLock(ctx); err != nil {
-		return err, false
-	}
-	defer func() {
-		if err := mutex.Unlock(ctx); err != nil {
-			n.logger.Errorf(ctx, "cannot unlock the lock: %s", err)
-		}
-	}()
-
-	// Purge job in bridge repository
-	if _, err = n.keboolaBridgeRepository.Job().Purge(&job).RequireLock(mutex).Do(ctx).ResultOrErr(); err != nil {
-		err = errors.PrefixErrorf(err, `cannot delete finished storage job "%s"`, job.JobKey)
-		n.logger.Error(ctx, err.Error())
-		return err, false
-	}
-
-	// Log file details
-	n.logger.Infof(ctx, `deleted finished storage job`)
-
-	return nil, true
-}