billing/worker/worker.go

package billingworker

import (
	"context"
	"fmt"
	"log/slog"
	"slices"
	"time"

	"github.com/ThreeDotsLabs/watermill/message"

	"github.com/openmeterio/openmeter/openmeter/billing"
	"github.com/openmeterio/openmeter/openmeter/billing/charges"
	chargesasyncadvance "github.com/openmeterio/openmeter/openmeter/billing/charges/worker/asyncadvance"
	"github.com/openmeterio/openmeter/openmeter/billing/worker/asyncadvance"
	billingworkercollect "github.com/openmeterio/openmeter/openmeter/billing/worker/collect"
	"github.com/openmeterio/openmeter/openmeter/billing/worker/subscriptionsync"
	"github.com/openmeterio/openmeter/openmeter/subscription"
	"github.com/openmeterio/openmeter/openmeter/watermill/eventbus"
	"github.com/openmeterio/openmeter/openmeter/watermill/grouphandler"
	"github.com/openmeterio/openmeter/openmeter/watermill/router"
)

type WorkerOptions struct {
	SystemEventsTopic string

	Router   router.Options
	EventBus eventbus.Publisher

	Logger *slog.Logger

	BillingService          billing.Service
	BillingSubscriptionSync subscriptionsync.Service
	BillingCollector        *billingworkercollect.InvoiceCollector
	// ChargesService is optional; when non-nil the worker handles AdvanceChargesEvent.
	ChargesService charges.ChargeService
	// External connectors

	SubscriptionService subscription.Service
	LockdownNamespaces  []string
}

func (w WorkerOptions) Validate() error {
	if w.SystemEventsTopic == "" {
		return fmt.Errorf("system events topic is required")
	}

	if err := w.Router.Validate(); err != nil {
		return fmt.Errorf("router: %w", err)
	}

	if w.EventBus == nil {
		return fmt.Errorf("event bus is required")
	}

	if w.Logger == nil {
		return fmt.Errorf("logger is required")
	}

	if w.BillingService == nil {
		return fmt.Errorf("billing service is required")
	}

	if w.SubscriptionService == nil {
		return fmt.Errorf("subscription service is required")
	}

	if w.BillingSubscriptionSync == nil {
		return fmt.Errorf("billing subscription sync handler is required")
	}

	if w.BillingCollector == nil {
		return fmt.Errorf("billing collector is required")
	}

	return nil
}

type Worker struct {
	router *message.Router

	billingService             billing.Service
	subscriptionSync           subscriptionsync.Service
	invoiceCollector           *billingworkercollect.InvoiceCollector
	asyncAdvanceHandler        *asyncadvance.Handler
	asyncAdvanceChargesHandler *chargesasyncadvance.Handler
	nonPublishingHandler       *grouphandler.NoPublishingHandler
	lockdownNamespaces         []string
}

func New(opts WorkerOptions) (*Worker, error) {
	if err := opts.Validate(); err != nil {
		return nil, err
	}

	// asyncAdvancer
	asyncAdvancer, err := asyncadvance.New(asyncadvance.Config{
		Logger:         opts.Logger,
		BillingService: opts.BillingService,
	})
	if err != nil {
		return nil, err
	}

	var asyncAdvanceChargesHandler *chargesasyncadvance.Handler
	if opts.ChargesService != nil {
		asyncAdvanceChargesHandler, err = chargesasyncadvance.New(chargesasyncadvance.Config{
			Logger:         opts.Logger,
			ChargesService: opts.ChargesService,
		})
		if err != nil {
			return nil, err
		}
	}

	worker := &Worker{
		billingService:             opts.BillingService,
		subscriptionSync:           opts.BillingSubscriptionSync,
		invoiceCollector:           opts.BillingCollector,
		asyncAdvanceHandler:        asyncAdvancer,
		asyncAdvanceChargesHandler: asyncAdvanceChargesHandler,
		lockdownNamespaces:         opts.LockdownNamespaces,
	}

	router, err := router.NewDefaultRouter(opts.Router)
	if err != nil {
		return nil, err
	}

	worker.router = router

	eventHandler, err := worker.eventHandler(opts)
	if err != nil {
		return nil, err
	}

	worker.nonPublishingHandler = eventHandler

	router.AddConsumerHandler(
		"billing_worker_system_events",
		opts.SystemEventsTopic,
		opts.Router.Subscriber,
		worker.nonPublishingHandler.Handle,
	)

	return worker, nil
}

// AddHandler adds an additional handler to the list of event handlers.
// Handlers are called in the order they are added and run after the built in handlers.
// In the case of any handler returning an error, the event will be retried so it is important that all handlers are idempotent.
func (w *Worker) AddHandler(handler grouphandler.GroupEventHandler) {
	w.nonPublishingHandler.AddHandler(handler)
}

func (w *Worker) eventHandler(opts WorkerOptions) (*grouphandler.NoPublishingHandler, error) {
	handler, err := grouphandler.NewNoPublishingHandler(
		opts.EventBus.Marshaler(),
		opts.Router.MetricMeter,

		grouphandler.NewGroupEventHandler(func(ctx context.Context, event *subscription.CreatedEvent) error {
			if event != nil && slices.Contains(w.lockdownNamespaces, event.Subscription.Namespace) {
				return nil
			}

			return w.subscriptionSync.SyncByViewAndInvoiceCustomer(ctx, event.SubscriptionView, time.Now())
		}),
		grouphandler.NewGroupEventHandler(func(ctx context.Context, event *subscription.CancelledEvent) error {
			if event != nil && slices.Contains(w.lockdownNamespaces, event.Subscription.Namespace) {
				return nil
			}

			return w.subscriptionSync.HandleCancelledEvent(ctx, event)
		}),
		grouphandler.NewGroupEventHandler(func(ctx context.Context, event *subscription.DeletedEvent) error {
			if event != nil && slices.Contains(w.lockdownNamespaces, event.Subscription.Namespace) {
				return nil
			}

			return w.subscriptionSync.HandleDeletedEvent(ctx, event)
		}),
		grouphandler.NewGroupEventHandler(func(ctx context.Context, event *subscription.ContinuedEvent) error {
			if event != nil && slices.Contains(w.lockdownNamespaces, event.Subscription.Namespace) {
				return nil
			}

			return w.subscriptionSync.SyncByViewAndInvoiceCustomer(ctx, event.SubscriptionView, time.Now())
		}),
		grouphandler.NewGroupEventHandler(func(ctx context.Context, event *subscription.UpdatedEvent) error {
			if event != nil && slices.Contains(w.lockdownNamespaces, event.UpdatedView.Subscription.Namespace) {
				return nil
			}

			return w.subscriptionSync.SyncByViewAndInvoiceCustomer(ctx, event.UpdatedView, time.Now())
		}),
		grouphandler.NewGroupEventHandler(func(ctx context.Context, event *subscription.SubscriptionSyncEvent) error {
			if event != nil && slices.Contains(w.lockdownNamespaces, event.Subscription.Namespace) {
				return nil
			}

			return w.subscriptionSync.HandleSubscriptionSyncEvent(ctx, event)
		}),
		grouphandler.NewGroupEventHandler(func(ctx context.Context, event *billing.AdvanceStandardInvoiceEvent) error {
			if event != nil && slices.Contains(w.lockdownNamespaces, event.Invoice.Namespace) {
				return nil
			}

			return w.asyncAdvanceHandler.Handle(ctx, event)
		}),
		grouphandler.NewGroupEventHandler(func(ctx context.Context, event *billing.CollectCustomerInvoicesEvent) error {
			if event != nil && slices.Contains(w.lockdownNamespaces, event.Namespace) {
				return nil
			}

			return w.invoiceCollector.HandleCollectCustomerInvoicesEvent(ctx, event)
		}),
		grouphandler.NewGroupEventHandler(func(ctx context.Context, event *charges.AdvanceChargesEvent) error {
			if w.asyncAdvanceChargesHandler == nil {
				return nil
			}

			if event != nil && slices.Contains(w.lockdownNamespaces, event.Namespace) {
				return nil
			}

			return w.asyncAdvanceChargesHandler.Handle(ctx, event)
		}),
		grouphandler.NewGroupEventHandler(func(ctx context.Context, event *billing.StandardInvoiceCreatedEvent) error {
			if event != nil && slices.Contains(w.lockdownNamespaces, event.Invoice.Namespace) {
				return nil
			}

			return w.subscriptionSync.HandleInvoiceCreation(ctx, event)
		}),
	)
	if err != nil {
		return nil, fmt.Errorf("failed to create event handler: %w", err)
	}

	return handler, nil
}

func (w *Worker) Run(ctx context.Context) error {
	return w.router.Run(ctx)
}

func (w *Worker) Close() error {
	return w.router.Close()
}