bbolt abstraction for statistic tracking

cmd/launcher/launcher.go

@@ -185,9 +185,14 @@ func runLauncher(ctx context.Context, cancel func(), multiSlogger, systemMultiSl
 		return fmt.Errorf("failed to create stores: %w", err)
 	}
 
+	statTracker, err := agentbbolt.NewStatStore(slogger, db)
+	if err != nil {
+		return fmt.Errorf("failed to create stat tracker: %w", err)
+	}
+
 	fcOpts := []flags.Option{flags.WithCmdLineOpts(opts)}
 	flagController := flags.NewFlagController(slogger, stores[storage.AgentFlagsStore], fcOpts...)
-	k := knapsack.New(stores, flagController, db, multiSlogger, systemMultiSlogger)
+	k := knapsack.New(stores, flagController, statTracker, multiSlogger, systemMultiSlogger)
 
 	go runOsqueryVersionCheck(ctx, slogger, k.LatestOsquerydPath(ctx))
 	go timemachine.AddExclusions(ctx, k)

ee/agent/knapsack/knapsack.go

@@ -14,7 +14,6 @@ import (
 	"github.com/kolide/launcher/ee/tuf"
 	"github.com/kolide/launcher/pkg/autoupdate"
 	"github.com/kolide/launcher/pkg/log/multislogger"
-	"go.etcd.io/bbolt"
 )
 
 // type alias Flags, so that we can embed it inside knapsack, as `flags` and not `Flags`
@@ -26,22 +25,15 @@ type knapsack struct {
 	stores map[storage.Store]types.KVStore
 	// Embed flags so we get all the flag interfaces
 	flags
-
-	// BboltDB is the underlying bbolt database.
-	// Ideally, we can eventually remove this. This is only here because some parts of the codebase
-	// like the osquery extension have a direct dependency on bbolt and need this reference.
-	// If we are able to abstract bbolt out completely in these areas, we should be able to
-	// remove this field and prevent "leaking" bbolt into places it doesn't need to.
-	db *bbolt.DB
-
+	storageStatTracker     types.StorageStatTracker
 	slogger, systemSlogger *multislogger.MultiSlogger
 
 	// This struct is a work in progress, and will be iteratively added to as needs arise.
 	// Some potential future additions include:
 	// Querier
 }
 
-func New(stores map[storage.Store]types.KVStore, flags types.Flags, db *bbolt.DB, slogger, systemSlogger *multislogger.MultiSlogger) *knapsack {
+func New(stores map[storage.Store]types.KVStore, flags types.Flags, sStatTracker types.StorageStatTracker, slogger, systemSlogger *multislogger.MultiSlogger) *knapsack {
 	if slogger == nil {
 		slogger = multislogger.New()
 	}
@@ -50,11 +42,11 @@ func New(stores map[storage.Store]types.KVStore, flags types.Flags, db *bbolt.DB
 	}
 
 	k := &knapsack{
-		db:            db,
-		flags:         flags,
-		stores:        stores,
-		slogger:       slogger,
-		systemSlogger: systemSlogger,
+		storageStatTracker: sStatTracker,
+		flags:              flags,
+		stores:             stores,
+		slogger:            slogger,
+		systemSlogger:      systemSlogger,
 	}
 
 	return k
@@ -74,9 +66,9 @@ func (k *knapsack) AddSlogHandler(handler ...slog.Handler) {
 	k.systemSlogger.AddHandler(handler...)
 }
 
-// BboltDB interface methods
-func (k *knapsack) BboltDB() *bbolt.DB {
-	return k.db
+// storage stat tracking interface methods
+func (k *knapsack) StorageStatTracker() types.StorageStatTracker {
+	return k.storageStatTracker
 }
 
 // Stores interface methods

ee/agent/storage/bbolt/stats_bbolt.go

@@ -0,0 +1,112 @@
+package agentbbolt
+
+import (
+	"encoding/json"
+	"fmt"
+	"log/slog"
+
+	"go.etcd.io/bbolt"
+)
+
+type bucketStatsHolder struct {
+	Stats        bbolt.BucketStats
+	FillPercent  float64
+	NumberOfKeys int
+	Size         int
+}
+
+type dbStatsHolder struct {
+	Stats bbolt.TxStats
+	Size  int64
+}
+
+type Stats struct {
+	DB      dbStatsHolder
+	Buckets map[string]bucketStatsHolder
+}
+
+type bboltStatStore struct {
+	slogger *slog.Logger
+	db      *bbolt.DB
+}
+
+// NewStatStore provides a wrapper around a bbolt.DB connection. This is done at the global (above bucket) level
+// and should be used for any operations regarding the collection of storage statistics
+func NewStatStore(slogger *slog.Logger, db *bbolt.DB) (*bboltStatStore, error) {
+	if db == nil {
+		return nil, NoDbError{}
+	}
+
+	s := &bboltStatStore{
+		slogger: slogger,
+		db:      db,
+	}
+
+	return s, nil
+}
+
+func (s *bboltStatStore) SizeBytes() (int64, error) {
+	if s == nil || s.db == nil {
+		return 0, NoDbError{}
+	}
+
+	var size int64
+
+	if err := s.db.View(func(tx *bbolt.Tx) error {
+		size = tx.Size()
+		return nil
+	}); err != nil {
+		return 0, fmt.Errorf("creating view tx to check size stat: %w", err)
+	}
+
+	return size, nil
+}
+
+// GetStats returns a json blob containing both global and bucket-level
+// statistics. Note that the bucketName set does not impact the output, all buckets
+// will be traversed for stats regardless
+func (s *bboltStatStore) GetStats() ([]byte, error) {
+	if s == nil || s.db == nil {
+		return nil, NoDbError{}
+	}
+
+	stats := &Stats{
+		Buckets: make(map[string]bucketStatsHolder),
+	}
+
+	if err := s.db.View(func(tx *bbolt.Tx) error {
+		stats.DB.Stats = tx.Stats()
+		stats.DB.Size = tx.Size()
+
+		if err := tx.ForEach(bucketStatsFunc(stats)); err != nil {
+			return fmt.Errorf("dumping bucket: %w", err)
+		}
+		return nil
+	}); err != nil {
+		return nil, fmt.Errorf("creating view tx: %w", err)
+	}
+
+	statsJson, err := json.Marshal(stats)
+	if err != nil {
+		return nil, err
+	}
+
+	return statsJson, nil
+}
+
+func bucketStatsFunc(stats *Stats) func([]byte, *bbolt.Bucket) error {
+	return func(name []byte, b *bbolt.Bucket) error {
+		bstats := b.Stats()
+
+		// KeyN is the number of keys
+		// LeafAlloc is pretty close the number of bytes used
+		stats.Buckets[string(name)] = bucketStatsHolder{
+			Stats:        bstats,
+			FillPercent:  b.FillPercent,
+			NumberOfKeys: bstats.KeyN,
+			Size:         bstats.LeafAlloc,
+		}
+
+		return nil
+	}
+}

ee/agent/types/keyvalue_store.go

@@ -55,6 +55,16 @@ type Appender interface {
 	AppendValues(values ...[]byte) error
 }
 
+// StorageStatTracker is an interface towards exposing the
+// statistics of an underlying storage methodology
+type StorageStatTracker interface {
+	// GetStats will return any relevant global and bucket-level statistics
+	// for the underlying storage methodology, expected as a json blob
+	GetStats() ([]byte, error)
+	// Size will return the total size in bytes for the stored key-value pairs
+	SizeBytes() (int64, error)
+}
+
 // GetterSetter is an interface that groups the Get and Set methods.
 type GetterSetter interface {
 	Getter

ee/agent/types/knapsack.go

@@ -6,7 +6,6 @@ import "context"
 // launcher code and are typically valid for the lifetime of the launcher application instance.
 type Knapsack interface {
 	Stores
-	BboltDB
 	Flags
 	Slogger
 	// LatestOsquerydPath finds the path to the latest osqueryd binary, after accounting for updates.
@@ -15,4 +14,5 @@ type Knapsack interface {
 	ReadEnrollSecret() (string, error)
 	// CurrentEnrollmentStatus returns the current enrollment status of the launcher installation
 	CurrentEnrollmentStatus() (EnrollmentStatus, error)
+	StorageStatTracker() StorageStatTracker
 }