logarchiver/internal/cli/run.go

package cli

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"os/signal"
	"syscall"
	"time"

	"git.unkin.net/unkin/logarchiver/internal/archiver"
	"git.unkin.net/unkin/logarchiver/internal/batcher"
	"git.unkin.net/unkin/logarchiver/internal/config"
	"git.unkin.net/unkin/logarchiver/internal/consumer"
	"git.unkin.net/unkin/logarchiver/internal/index"
	"git.unkin.net/unkin/logarchiver/internal/metrics"
	"git.unkin.net/unkin/logarchiver/internal/s3store"
	"git.unkin.net/unkin/logarchiver/internal/vaultgpg"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/spf13/cobra"
)

func newRunCmd() *cobra.Command {
	return &cobra.Command{
		Use:   "run",
		Short: "Run the archiver service (JetStream consumer -> S3 + index)",
		RunE: func(cmd *cobra.Command, _ []string) error {
			cfg, err := loadConfig()
			if err != nil {
				return err
			}
			return runService(cmd.Context(), cfg)
		},
	}
}

// persistAdapter bridges *archiver.Archiver to consumer.Persister (different
// StoreResult types across package boundaries).
type persistAdapter struct{ a *archiver.Archiver }

func (p persistAdapter) Store(ctx context.Context, b *batcher.Batch) (consumer.StoreResult, error) {
	res, err := p.a.Store(ctx, b)
	return consumer.StoreResult{
		ObjectKey:   res.ObjectKey,
		Events:      res.Events,
		RawBytes:    res.RawBytes,
		StoredBytes: res.StoredBytes,
	}, err
}

func runService(parent context.Context, cfg config.Config) error {
	log := newLogger(cfg.Log)
	slog.SetDefault(log)

	ctx, stop := signal.NotifyContext(parent, syscall.SIGINT, syscall.SIGTERM)
	defer stop()

	// Metrics.
	var met *metrics.Metrics
	var metricsSrv *http.Server
	if cfg.Metrics.Enabled {
		reg := prometheus.NewRegistry()
		met = metrics.New(reg)
		mux := http.NewServeMux()
		mux.Handle("/metrics", promhttp.HandlerFor(reg, promhttp.HandlerOpts{}))
		mux.HandleFunc("/healthz", func(w http.ResponseWriter, _ *http.Request) {
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte("ok"))
		})
		metricsSrv = &http.Server{Addr: cfg.Metrics.Address, Handler: mux, ReadHeaderTimeout: 5 * time.Second}
		go func() {
			log.Info("metrics listening", "addr", cfg.Metrics.Address)
			if err := metricsSrv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
				log.Error("metrics server failed", "err", err)
			}
		}()
	}

	// Vault client only when we source the public key from Vault; the service
	// never needs Vault otherwise (file-mounted pubkey is the default).
	var vc *vaultgpg.Client
	if cfg.Crypto.Source == config.PubkeyVault {
		var err error
		vc, err = vaultgpg.New(ctx, vaultConfig(cfg.Crypto.Vault))
		if err != nil {
			return fmt.Errorf("vault init: %w", err)
		}
	}

	// Public key provider.
	loader, err := archiver.PubkeyLoaderFromConfig(cfg.Crypto, vc)
	if err != nil {
		return err
	}
	pubkeys, err := archiver.NewPubkeyProvider(ctx, loader)
	if err != nil {
		return fmt.Errorf("load public key: %w", err)
	}
	log.Info("public key loaded",
		"source", cfg.Crypto.Source, "key_name", cfg.Crypto.KeyName,
		"fingerprint", pubkeys.Current().Fingerprint)
	go refreshPubkey(ctx, log, pubkeys, cfg.Crypto.RefreshInterval)

	// S3.
	store, err := s3store.New(ctx, s3store.Config{
		Endpoint:  cfg.S3.Endpoint,
		Bucket:    cfg.S3.Bucket,
		Region:    cfg.S3.Region,
		PathStyle: cfg.S3.PathStyle,
		CAFile:    cfg.S3.CAFile,
	})
	if err != nil {
		return fmt.Errorf("s3 init: %w", err)
	}

	// Index.
	var idx index.Index
	if cfg.Index.Enabled {
		ch, err := index.NewClickHouse(ctx, indexConfig(cfg.Index))
		if err != nil {
			return fmt.Errorf("clickhouse init: %w", err)
		}
		defer func() { _ = ch.Close() }()
		idx = ch
	}

	keys, err := archiver.NewKeyBuilder(cfg.S3.KeyPrefix)
	if err != nil {
		return err
	}

	arch, err := archiver.New(archiver.Options{
		Keys:      keys,
		Pubkeys:   pubkeys,
		Store:     store,
		Index:     idx,
		KeyName:   cfg.Crypto.KeyName,
		FrameSize: cfg.Crypto.FrameSize,
		Metrics:   met,
	})
	if err != nil {
		return err
	}

	// NATS + consumer.
	nc, js, err := consumer.Connect(cfg.NATS)
	if err != nil {
		return err
	}
	defer nc.Close()
	cons, err := consumer.EnsureConsumer(ctx, js, cfg.NATS)
	if err != nil {
		return err
	}
	log.Info("consumer bound",
		"stream", cfg.NATS.Stream, "durable", cfg.NATS.Durable, "subjects", cfg.NATS.Subjects)

	bat := batcher.New(batcher.Limits{
		MaxBytes:  cfg.Batch.MaxBytes,
		MaxEvents: cfg.Batch.MaxEvents,
		MaxAge:    cfg.Batch.MaxAge,
	})

	var runnerMetrics consumer.Metrics
	if met != nil {
		runnerMetrics = met
	}
	runner := consumer.NewRunner(consumer.Options{
		Consumer:   cons,
		Batcher:    bat,
		Persister:  persistAdapter{a: arch},
		Logger:     log,
		Metrics:    runnerMetrics,
		FetchBatch: cfg.NATS.FetchBatch,
		PollWait:   pollWait(cfg.Batch.MaxAge),
	})

	runErr := runner.Run(ctx)

	if metricsSrv != nil {
		shCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		_ = metricsSrv.Shutdown(shCtx)
		cancel()
	}
	if runErr != nil && !errors.Is(runErr, context.Canceled) {
		return runErr
	}
	log.Info("shutdown complete")
	return nil
}

// pollWait picks a fetch/age-check interval that is a fraction of MaxAge so
// aged batches flush promptly, clamped to a sane range.
func pollWait(maxAge time.Duration) time.Duration {
	if maxAge <= 0 {
		return time.Second
	}
	w := maxAge / 10
	if w < time.Second {
		w = time.Second
	}
	if w > 10*time.Second {
		w = 10 * time.Second
	}
	return w
}

func refreshPubkey(ctx context.Context, log *slog.Logger, p *archiver.PubkeyProvider, every time.Duration) {
	if every <= 0 {
		return
	}
	t := time.NewTicker(every)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			if err := p.Refresh(ctx); err != nil {
				log.Warn("pubkey refresh failed; keeping previous key", "err", err)
				continue
			}
			log.Debug("pubkey refreshed", "fingerprint", p.Current().Fingerprint)
		}
	}
}

func vaultConfig(v config.VaultConfig) vaultgpg.Config {
	return vaultgpg.Config{
		Address:    v.Address,
		Mount:      v.Mount,
		AuthMethod: v.AuthMethod,
		K8sRole:    v.K8sRole,
		K8sMount:   v.K8sMount,
		K8sJWTPath: v.K8sJWTPath,
		CAFile:     v.CAFile,
	}
}

func indexConfig(i config.IndexConfig) index.Config {
	return index.Config{
		Address:  i.Address,
		Database: i.Database,
		Table:    i.Table,
		Username: i.Username,
		Password: i.Password,
		TLS:      i.TLS,
	}
}