[usage] Refactor controller package into scheduler

components/usage/pkg/controller/reconciler.go → components/usage/pkg/scheduler/job.go

@@ -2,55 +2,76 @@
 // Licensed under the GNU Affero General Public License (AGPL).
 // See License-AGPL.txt in the project root for license information.
 
-package controller
+package scheduler
 
 import (
 	"context"
 	"fmt"
 	"github.com/gitpod-io/gitpod/common-go/log"
 	v1 "github.com/gitpod-io/gitpod/usage-api/v1"
+	"github.com/robfig/cron"
 	"google.golang.org/protobuf/types/known/timestamppb"
 	"time"
 )
 
-type Reconciler interface {
-	Reconcile() error
+type Job interface {
+	Run() error
 }
 
-type ReconcilerFunc func() error
+func NewLedgerTriggerJobSpec(schedule time.Duration, job Job) (JobSpec, error) {
+	parsed, err := cron.Parse(fmt.Sprintf("@every %s", schedule.String()))
+	if err != nil {
+		return JobSpec{}, fmt.Errorf("failed to parse ledger job schedule: %w", err)
+	}
 
-func (f ReconcilerFunc) Reconcile() error {
-	return f()
+	return JobSpec{
+		Job:      job,
+		ID:       "ledger",
+		Schedule: parsed,
+	}, nil
 }
 
-func NewLedgerReconciler(usageClient v1.UsageServiceClient, billingClient v1.BillingServiceClient) *LedgerReconciler {
-	return &LedgerReconciler{
+func NewLedgerTrigger(usageClient v1.UsageServiceClient, billingClient v1.BillingServiceClient) *LedgerJob {
+	return &LedgerJob{
 		usageClient:   usageClient,
 		billingClient: billingClient,
+
+		running: make(chan struct{}, 1),
 	}
 }
 
-type LedgerReconciler struct {
+type LedgerJob struct {
 	usageClient   v1.UsageServiceClient
 	billingClient v1.BillingServiceClient
+
+	running chan struct{}
 }
 
-func (r *LedgerReconciler) Reconcile() (err error) {
+func (r *LedgerJob) Run() (err error) {
 	ctx := context.Background()
 
+	select {
+	// attempt a write to signal we want to run
+	case r.running <- struct{}{}:
+		// we managed to write, there's no other job executing. Cases are not fall through so we continue executing our main logic.
+		defer func() {
+			// signal job completed
+			<-r.running
+		}()
+	default:
+		// we could not write, so another instance is already running. Skip current run.
+		log.Infof("Skipping ledger run, another run is already in progress.")
+		return nil
+	}
+
 	now := time.Now().UTC()
 	hourAgo := now.Add(-1 * time.Hour)
 
-	reportUsageReconcileStarted()
-	defer func() {
-		reportUsageReconcileFinished(time.Since(now), err)
-	}()
-
 	logger := log.
 		WithField("from", hourAgo).
 		WithField("to", now)
 
-	logger.Info("Starting ledger reconciliation.")
+	logger.Info("Running ledger job. Reconciling usage records.")
 	_, err = r.usageClient.ReconcileUsageWithLedger(ctx, &v1.ReconcileUsageWithLedgerRequest{
 		From: timestamppb.New(hourAgo),
 		To:   timestamppb.New(now),

components/usage/pkg/scheduler/job_test.go

@@ -0,0 +1,73 @@
+// Copyright (c) 2022 Gitpod GmbH. All rights reserved.
+// Licensed under the GNU Affero General Public License (AGPL).
+// See License-AGPL.txt in the project root for license information.
+
+package scheduler
+
+import (
+	"context"
+	v1 "github.com/gitpod-io/gitpod/usage-api/v1"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func TestLedgerJob_PreventsConcurrentInvocations(t *testing.T) {
+	client := &fakeUsageClient{}
+	job := NewLedgerTrigger(client, nil)
+
+	invocations := 3
+	wg := sync.WaitGroup{}
+	wg.Add(invocations)
+	for i := 0; i < invocations; i++ {
+		go func() {
+			_ = job.Run()
+			wg.Done()
+		}()
+	}
+	wg.Wait()
+
+	require.Equal(t, 1, int(client.ReconcileUsageWithLedgerCallCount))
+}
+
+func TestLedgerJob_CanRunRepeatedly(t *testing.T) {
+	client := &fakeUsageClient{}
+	job := NewLedgerTrigger(client, nil)
+
+	_ = job.Run()
+	_ = job.Run()
+
+	require.Equal(t, 2, int(client.ReconcileUsageWithLedgerCallCount))
+}
+
+type fakeUsageClient struct {
+	ReconcileUsageWithLedgerCallCount int32
+}
+
+// GetCostCenter retrieves the active cost center for the given attributionID
+func (c *fakeUsageClient) GetCostCenter(ctx context.Context, in *v1.GetCostCenterRequest, opts ...grpc.CallOption) (*v1.GetCostCenterResponse, error) {
+	return nil, status.Error(codes.Unauthenticated, "not implemented")
+}
+
+// SetCostCenter stores the given cost center
+func (c *fakeUsageClient) SetCostCenter(ctx context.Context, in *v1.SetCostCenterRequest, opts ...grpc.CallOption) (*v1.SetCostCenterResponse, error) {
+	return nil, status.Error(codes.Unauthenticated, "not implemented")
+}
+
+// Triggers reconciliation of usage with ledger implementation.
+func (c *fakeUsageClient) ReconcileUsageWithLedger(ctx context.Context, in *v1.ReconcileUsageWithLedgerRequest, opts ...grpc.CallOption) (*v1.ReconcileUsageWithLedgerResponse, error) {
+	atomic.AddInt32(&c.ReconcileUsageWithLedgerCallCount, 1)
+	time.Sleep(50 * time.Millisecond)
+
+	return nil, status.Error(codes.Unauthenticated, "not implemented")
+}
+
+// ListUsage retrieves all usage for the specified attributionId and theb given time range
+func (c *fakeUsageClient) ListUsage(ctx context.Context, in *v1.ListUsageRequest, opts ...grpc.CallOption) (*v1.ListUsageResponse, error) {
+	return nil, status.Error(codes.Unauthenticated, "not implemented")
+}

components/usage/pkg/controller/reporter.go → components/usage/pkg/scheduler/reporter.go

@@ -2,7 +2,7 @@
 // Licensed under the GNU Affero General Public License (AGPL).
 // See License-AGPL.txt in the project root for license information.
 
-package controller
+package scheduler
 
 import (
 	"fmt"
@@ -16,26 +16,26 @@ const (
 )
 
 var (
-	reconcileStartedTotal = prometheus.NewCounterVec(prometheus.CounterOpts{
+	jobStartedSeconds = prometheus.NewCounterVec(prometheus.CounterOpts{
 		Namespace: namespace,
 		Subsystem: subsystem,
-		Name:      "reconcile_started_total",
-		Help:      "Number of usage reconciliation runs started",
-	}, []string{})
+		Name:      "scheduler_job_started_total",
+		Help:      "Number of jobs started",
+	}, []string{"job"})
 
-	reconcileStartedDurationSeconds = prometheus.NewHistogramVec(prometheus.HistogramOpts{
+	jobCompletedSeconds = prometheus.NewHistogramVec(prometheus.HistogramOpts{
 		Namespace: namespace,
 		Subsystem: subsystem,
-		Name:      "reconcile_completed_duration_seconds",
-		Help:      "Histogram of reconcile duration",
+		Name:      "scheduler_job_completed_seconds",
+		Help:      "Histogram of job duration",
 		Buckets:   prometheus.LinearBuckets(30, 30, 10), // every 30 secs, starting at 30secs
-	}, []string{"outcome"})
+	}, []string{"job", "outcome"})
 )
 
 func RegisterMetrics(reg *prometheus.Registry) error {
 	metrics := []prometheus.Collector{
-		reconcileStartedTotal,
-		reconcileStartedDurationSeconds,
+		jobStartedSeconds,
+		jobCompletedSeconds,
 	}
 	for _, metric := range metrics {
 		err := reg.Register(metric)
@@ -47,14 +47,14 @@ func RegisterMetrics(reg *prometheus.Registry) error {
 	return nil
 }
 
-func reportUsageReconcileStarted() {
-	reconcileStartedTotal.WithLabelValues().Inc()
+func reportJobStarted(id string) {
+	jobStartedSeconds.WithLabelValues(id).Inc()
 }
 
-func reportUsageReconcileFinished(duration time.Duration, err error) {
+func reportJobCompleted(id string, duration time.Duration, err error) {
 	outcome := "success"
 	if err != nil {
 		outcome = "error"
 	}
-	reconcileStartedDurationSeconds.WithLabelValues(outcome).Observe(duration.Seconds())
+	jobCompletedSeconds.WithLabelValues(id, outcome).Observe(duration.Seconds())
 }