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"
	"database/sql"
	"errors"
	"fmt"
	"path/filepath"
	"sync"
	"time"

	"git.eeqj.de/sneak/routewatch/internal/cachedb"
	"git.eeqj.de/sneak/routewatch/internal/config"
	"git.eeqj.de/sneak/routewatch/internal/database"
	"git.eeqj.de/sneak/routewatch/internal/geo"
	"git.eeqj.de/sneak/routewatch/internal/iplist"
	"git.eeqj.de/sneak/routewatch/internal/logger"
	"git.eeqj.de/sneak/routewatch/internal/metrics"
	"git.eeqj.de/sneak/routewatch/internal/ptr"
	"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
	ipLists        *iplist.IPLists
	cacheDB        *sql.DB
	ptrCache       *ptr.Cache
	geoLocator     *geo.Locator
}

// ipListsDirName is the directory in the state directory that holds the IP
// list copies.
const ipListsDirName = "iplists"

// 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 IP list copies from disk and start downloading the lists.
	// The HTTP handlers read the server's IP lists without a lock, so they
	// are set before the handlers run.
	ipLists := iplist.New(iplist.Params{
		Dir:    filepath.Join(rw.config.StateDir, ipListsDirName),
		Lists:  iplist.Lists(),
		Logger: rw.logger.Logger,
	})
	ipLists.Start(ctx)
	rw.ipLists = ipLists
	rw.server.SetIPLists(ipLists)

	err = rw.startPTRCache(ctx)
	if err != nil {
		return err
	}

	err = rw.startGeoLocator(ctx)
	if err != nil {
		return err
	}

	// 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
// IP list downloads, the HTTP server, the geolocation cache and the PTR cache,
// 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.ipLists != nil {
		rw.ipLists.Stop(ctx)
	}

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

	if rw.geoLocator != nil {
		rw.geoLocator.Stop(ctx)
	}

	rw.stopPTRCache(ctx)

	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),
	)
}

// startPTRCache opens cache.db, loads the PTR cache from it and gives the
// server the cache. The HTTP handlers read it without a lock, so it is set
// before they run.
func (rw *RouteWatch) startPTRCache(ctx context.Context) error {
	cacheDB, err := cachedb.Open(ctx, rw.config.StateDir)
	if err != nil {
		return err
	}

	ptrCache := ptr.New(ptr.Params{
		DB:        cacheDB,
		Resolvers: ptr.Resolvers(),
		Logger:    rw.logger.Logger,
	})

	err = ptrCache.Start(ctx)
	if err != nil {
		_ = cacheDB.Close()

		return err
	}

	rw.cacheDB = cacheDB
	rw.ptrCache = ptrCache
	rw.server.SetPTRResolver(ptrCache)

	return nil
}

// startGeoLocator loads the geolocation cache from cache.db, which
// startPTRCache has opened, and gives the server the locator. The HTTP handlers
// read it without a lock, so it is set before they run.
func (rw *RouteWatch) startGeoLocator(ctx context.Context) error {
	locator, err := geo.New(geo.Params{
		DB:                rw.cacheDB,
		Dir:               rw.config.StateDir,
		GeoJSURL:          geo.GeoJSEndpoint,
		DatabaseURLFormat: geo.DatabaseDownloadFormat,
		Logger:            rw.logger.Logger,
	})
	if err != nil {
		return err
	}

	err = locator.Start(ctx)
	if err != nil {
		return err
	}

	rw.geoLocator = locator
	rw.server.SetGeoLocator(locator)

	return nil
}

// stopPTRCache stops the PTR cache, which writes the answers cache.db does not
// have yet, and closes cache.db, unless ctx has ended and a write may still be
// running.
func (rw *RouteWatch) stopPTRCache(ctx context.Context) {
	if rw.ptrCache == nil {
		return
	}

	rw.ptrCache.Stop(ctx)

	if ctx.Err() != nil {
		return
	}

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

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