routewatch/internal/routewatch/app.go

// Package routewatch contains the primary RouteWatch type that represents a
// running instance of the application and contains pointers to its core
// dependencies, and is responsible for initialization. It also contains the
// ASN WHOIS fetcher for background updates and the database maintainer for
// background maintenance tasks.
package routewatch

import (
	"context"
	"errors"
	"sync"
	"time"

	"git.eeqj.de/sneak/routewatch/internal/config"
	"git.eeqj.de/sneak/routewatch/internal/database"
	"git.eeqj.de/sneak/routewatch/internal/logger"
	"git.eeqj.de/sneak/routewatch/internal/metrics"
	"git.eeqj.de/sneak/routewatch/internal/server"
	"git.eeqj.de/sneak/routewatch/internal/streamer"

	"go.uber.org/fx"
)

// errNonBatchedNotImplemented is returned by Run when batched database writes
// are disabled, since only the batched handlers exist.
var errNonBatchedNotImplemented = errors.New("non-batched handlers not implemented")

// Dependencies contains all dependencies for RouteWatch
type Dependencies struct {
	fx.In

	DB       database.Store
	Streamer *streamer.Streamer
	Server   *server.Server
	Logger   *logger.Logger
	Config   *config.Config
}

// RouteWatch represents the main application instance
type RouteWatch struct {
	db             database.Store
	streamer       *streamer.Streamer
	server         *server.Server
	logger         *logger.Logger
	maxRuntime     time.Duration
	shutdown       bool
	mu             sync.Mutex
	config         *config.Config
	dbHandler      *ASHandler
	peerHandler    *PeerHandler
	prefixHandler  *PrefixHandler
	peeringHandler *PeeringHandler
	asnFetcher     *ASNFetcher
	dbMaintainer   *DBMaintainer
}

// New creates a new RouteWatch instance
func New(deps Dependencies) *RouteWatch {
	rw := &RouteWatch{
		db:         deps.DB,
		streamer:   deps.Streamer,
		server:     deps.Server,
		logger:     deps.Logger,
		maxRuntime: deps.Config.MaxRuntime,
		config:     deps.Config,
	}

	return rw
}

// Run starts the RouteWatch application
func (rw *RouteWatch) Run(ctx context.Context) error {
	rw.logger.Info("Starting RouteWatch")

	// Apply runtime limit if specified
	if rw.maxRuntime > 0 {
		var cancel context.CancelFunc

		ctx, cancel = context.WithTimeout(ctx, rw.maxRuntime)
		defer cancel()

		rw.logger.Info("Running with time limit", "max_runtime", rw.maxRuntime)
	}

	// Register database handler to process BGP UPDATE messages
	if rw.config.EnableBatchedDatabaseWrites {
		rw.logger.Info("Using batched database handlers for improved performance")
		// ASHandler maintains the asns table
		rw.dbHandler = NewASHandler(rw.db, rw.logger)
		rw.streamer.RegisterHandler(rw.dbHandler)

		// PeerHandler maintains the bgp_peers table
		rw.peerHandler = NewPeerHandler(rw.db, rw.logger)
		rw.streamer.RegisterHandler(rw.peerHandler)

		// PrefixHandler maintains the prefixes and live_routes tables
		rw.prefixHandler = NewPrefixHandler(rw.db, rw.logger)
		rw.prefixHandler.SetMetricsTracker(rw.streamer.GetMetricsTracker())
		rw.streamer.RegisterHandler(rw.prefixHandler)

		// PeeringHandler maintains the asn_peerings table
		rw.peeringHandler = NewPeeringHandler(rw.db, rw.logger)
		rw.streamer.RegisterHandler(rw.peeringHandler)
	} else {
		// Non-batched handlers not implemented yet
		rw.logger.Error("Non-batched handlers not implemented")

		return errNonBatchedNotImplemented
	}

	// No longer need routing table handler - PrefixHandler maintains live_routes table

	// Start streaming
	err := rw.streamer.Start(ctx)
	if err != nil {
		return err
	}

	// Start HTTP server
	err = rw.server.Start(ctx)
	if err != nil {
		return err
	}

	// Start ASN WHOIS fetcher for background updates
	rw.asnFetcher = NewASNFetcher(rw.db, rw.logger.Logger)
	rw.asnFetcher.Start(ctx)
	rw.server.SetASNFetcher(rw.asnFetcher)

	// Start database maintenance goroutine
	rw.dbMaintainer = NewDBMaintainer(rw.db, rw.logger.Logger)
	rw.dbMaintainer.Start(ctx)

	// Wait for context cancellation
	<-ctx.Done()

	return nil
}

// Shutdown stops all services in order: the streamer, the handlers, which
// write what they hold, the WHOIS fetcher, the database maintainer and the
// HTTP server, and last the database, which it checkpoints and closes. Each
// stage waits at most until ctx ends, so Shutdown returns soon after: a
// handler logs what it could not write by then, and the HTTP server closes
// the connections of requests still running.
func (rw *RouteWatch) Shutdown(ctx context.Context) {
	rw.mu.Lock()
	if rw.shutdown {
		rw.mu.Unlock()

		return
	}

	rw.shutdown = true
	rw.mu.Unlock()

	// Stop the streamer first: it returns once the handlers have been handed
	// every message it queued, or when ctx ends. Then each handler writes
	// what it holds.
	rw.streamer.Stop(ctx)

	if rw.dbHandler != nil {
		rw.logger.Info("Flushing database handler")
		rw.dbHandler.Stop(ctx)
	}

	if rw.peerHandler != nil {
		rw.logger.Info("Flushing peer handler")
		rw.peerHandler.Stop(ctx)
	}

	if rw.prefixHandler != nil {
		rw.logger.Info("Flushing prefix handler")
		rw.prefixHandler.Stop(ctx)
	}

	if rw.peeringHandler != nil {
		rw.logger.Info("Flushing peering handler")
		rw.peeringHandler.Stop(ctx)
	}

	// Stop ASN WHOIS fetcher
	if rw.asnFetcher != nil {
		rw.asnFetcher.Stop(ctx)
	}

	// Stop database maintainer
	if rw.dbMaintainer != nil {
		rw.dbMaintainer.Stop(ctx)
	}

	err := rw.server.Stop(ctx)
	if err != nil {
		rw.logger.Error("Stopped the HTTP server before every request finished",
			"error", err)
	}

	rw.closeDatabase(ctx)

	// Log final metrics
	metrics := rw.streamer.GetMetrics()
	rw.logger.Info("Final metrics",
		"total_messages", metrics.TotalMessages,
		"total_bytes", metrics.TotalBytes,
		"messages_per_sec", metrics.MessagesPerSec,
		"bits_per_sec", metrics.BitsPerSec,
		"duration", time.Since(metrics.ConnectedSince),
	)
}

// closeDatabase checkpoints and closes the database, unless ctx has ended: a
// write may then still be running, and both would wait for it. SQLite rolls
// back a write cut off by the exit when the database next opens.
func (rw *RouteWatch) closeDatabase(ctx context.Context) {
	if ctx.Err() != nil {
		rw.logger.Error("Stop deadline passed; the database is not closed")

		return
	}

	err := rw.db.Checkpoint(ctx)
	if err != nil {
		rw.logger.Error("Failed to checkpoint the database", "error", err)
	}

	err = rw.db.Close()
	if err != nil {
		rw.logger.Error("Failed to close the database", "error", err)
	}
}

// getModule provides all fx dependencies
func getModule() []fx.Option {
	return []fx.Option{
		fx.Provide(
			logger.New,
			config.New,
			metrics.New,
			fx.Annotate(
				database.New,
				fx.As(new(database.Store)),
			),
			streamer.New,
			server.New,
			New,
		),
	}
}