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. The package also holds the handlers that
// write the feed's updates to the database in batches, and the flush loop
// that the AS, peer and prefix handlers share.
package routewatch

import (
	"context"
	"errors"
	"fmt"
	"path/filepath"
	"sync"
	"time"

	"git.eeqj.de/sneak/routewatch/internal/blocklist"
	"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/routetable"
	"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. Table is the routing
// table read from the database at startup, which the server reads too.
type Dependencies struct {
	fx.In

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

// RouteWatch represents the main application instance
type RouteWatch struct {
	db             database.Store
	table          *routetable.Table
	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
	blocklists     *blocklist.Blocklists
}

// blocklistsDirName is the directory in the state directory that holds the
// blocklist copies.
const blocklistsDirName = "blocklists"

// New creates a new RouteWatch instance
func New(deps Dependencies) *RouteWatch {
	rw := &RouteWatch{
		db:         deps.DB,
		table:      deps.Table,
		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 keeps the routing table in memory current, and writes
		// its changes to the routes table and new prefixes to the prefix
		// tables
		rw.prefixHandler = NewPrefixHandler(rw.db, rw.logger, rw.table)
		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
	}

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

	// Start ASN WHOIS fetcher for background updates. The HTTP handlers read
	// the server's fetcher without a lock, so it is set before they run.
	rw.asnFetcher = NewASNFetcher(rw.db, rw.logger.Logger)
	rw.asnFetcher.Start(ctx)
	rw.server.SetASNFetcher(rw.asnFetcher)

	// Load the blocklist copies from disk and start downloading the lists.
	// The HTTP handlers read the server's blocklists without a lock, so they
	// are set before the handlers run.
	blocklists := blocklist.New(blocklist.Params{
		Dir:    filepath.Join(rw.config.StateDir, blocklistsDirName),
		Lists:  blocklist.Lists(),
		Logger: rw.logger.Logger,
	})
	blocklists.Start(ctx)
	rw.blocklists = blocklists
	rw.server.SetBlocklists(blocklists)

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

	// 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, the
// blocklist downloads 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)
	}

	if rw.blocklists != nil {
		rw.blocklists.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)
	}
}

// readRouteTable reads the routing table from the database and logs how many
// prefix and origin AS pairs and peers it holds and how long the read took.
// fx calls it once, to build the server and RouteWatch, before the streamer
// connects, so the prefix handler applies the feed's updates to all of it. No
// start timeout limits the read: fx builds them before it runs a start hook.
func readRouteTable(
	db database.Store, logger *logger.Logger,
) (*routetable.Table, error) {
	start := time.Now()

	table, err := db.ReadRouteTable(context.Background())
	if err != nil {
		return nil, fmt.Errorf("failed to read the routing table: %w", err)
	}

	counts := table.Counts()
	logger.Info("Read the routing table", "pairs", counts.Pairs,
		"peers", counts.Peers, "duration_ms", time.Since(start).Milliseconds())

	return table, nil
}

// 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)),
			),
			readRouteTable,
			streamer.New,
			server.New,
			New,
		),
	}
}