Add additional spans for bbolt.MakeStores

kolide · Jan 24, 2025 · 0da640e · 0da640e
1 parent af13166
commit 0da640e
Show file tree

Hide file tree

Showing 6 changed files with 14 additions and 7 deletions.
diff --git a/ee/agent/storage/bbolt/keyvalue_store_bbolt.go b/ee/agent/storage/bbolt/keyvalue_store_bbolt.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"log/slog"
 
+	"github.com/kolide/launcher/pkg/traces"
 	"go.etcd.io/bbolt"
 )
 
@@ -34,7 +35,10 @@ type bboltKeyValueStore struct {
 	bucketName string
 }
 
-func NewStore(slogger *slog.Logger, db *bbolt.DB, bucketName string) (*bboltKeyValueStore, error) {
+func NewStore(ctx context.Context, slogger *slog.Logger, db *bbolt.DB, bucketName string) (*bboltKeyValueStore, error) {
+	_, span := traces.StartSpan(ctx, "bucket_name", bucketName)
+	defer span.End()
+
 	err := db.Update(func(tx *bbolt.Tx) error {
 		_, err := tx.CreateBucketIfNotExists([]byte(bucketName))
 		if err != nil {

diff --git a/ee/agent/storage/bbolt/stores_bbolt.go b/ee/agent/storage/bbolt/stores_bbolt.go
@@ -13,7 +13,7 @@ import (
 
 // MakeStores creates all the KVStores used by launcher
 func MakeStores(ctx context.Context, slogger *slog.Logger, db *bbolt.DB) (map[storage.Store]types.KVStore, error) {
-	_, span := traces.StartSpan(ctx)
+	ctx, span := traces.StartSpan(ctx)
 	defer span.End()
 
 	stores := make(map[storage.Store]types.KVStore)
@@ -37,7 +37,7 @@ func MakeStores(ctx context.Context, slogger *slog.Logger, db *bbolt.DB) (map[st
 	}
 
 	for _, storeName := range storeNames {
-		store, err := NewStore(slogger, db, storeName.String())
+		store, err := NewStore(ctx, slogger, db, storeName.String())
 		if err != nil {
 			return nil, fmt.Errorf("failed to create '%s' KVStore: %w", storeName, err)
 		}

diff --git a/ee/agent/storage/ci/keyvalue_store_ci.go b/ee/agent/storage/ci/keyvalue_store_ci.go
@@ -1,6 +1,7 @@
 package storageci
 
 import (
+	"context"
 	"fmt"
 	"log/slog"
 	"os"
@@ -23,7 +24,7 @@ func NewStore(t *testing.T, slogger *slog.Logger, bucketName string) (types.KVSt
 		return inmemory.NewStore(), nil
 	}
 
-	return agentbbolt.NewStore(slogger, SetupDB(t), bucketName)
+	return agentbbolt.NewStore(context.TODO(), slogger, SetupDB(t), bucketName)
 }
 
 // SetupDB is used for creating bbolt databases for testing

diff --git a/ee/agent/storage/ci/keyvalue_store_test.go b/ee/agent/storage/ci/keyvalue_store_test.go
@@ -1,6 +1,7 @@
 package storageci
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"sync"
@@ -16,7 +17,7 @@ import (
 
 func getStores(t *testing.T) []types.KVStore {
 	db := SetupDB(t)
-	bboltStore, err := agentbbolt.NewStore(multislogger.NewNopLogger(), db, "test_bucket")
+	bboltStore, err := agentbbolt.NewStore(context.TODO(), multislogger.NewNopLogger(), db, "test_bucket")
 	require.NoError(t, err)
 
 	stores := []types.KVStore{

diff --git a/ee/agent/storage/ci/stores_ci.go b/ee/agent/storage/ci/stores_ci.go
@@ -1,6 +1,7 @@
 package storageci
 
 import (
+	"context"
 	"fmt"
 	"log/slog"
 	"os"
@@ -54,7 +55,7 @@ func makeBboltStores(t *testing.T, slogger *slog.Logger, db *bbolt.DB, storeName
 	stores := make(map[storage.Store]types.KVStore)
 
 	for _, storeName := range storeNames {
-		store, err := agentbbolt.NewStore(slogger, db, storeName.String())
+		store, err := agentbbolt.NewStore(context.TODO(), slogger, db, storeName.String())
 		if err != nil {
 			return nil, fmt.Errorf("failed to create '%s' KVStore: %w", storeName, err)
 		}

diff --git a/pkg/osquery/extension_test.go b/pkg/osquery/extension_test.go
@@ -99,7 +99,7 @@ func TestNewExtensionDatabaseError(t *testing.T) {
 	}()
 
 	m := mocks.NewKnapsack(t)
-	m.On("ConfigStore").Return(agentbbolt.NewStore(multislogger.NewNopLogger(), db, storage.ConfigStore.String()))
+	m.On("ConfigStore").Return(agentbbolt.NewStore(context.TODO(), multislogger.NewNopLogger(), db, storage.ConfigStore.String()))
 	m.On("Slogger").Return(multislogger.NewNopLogger()).Maybe()
 
 	e, err := NewExtension(context.TODO(), &mock.KolideService{}, m, ulid.New(), ExtensionOpts{})