[jaeger-v2] Add badger e2e integration test (jaegertracing#5355)

## Which problem is this PR solving?
- part of jaegertracing#5254
## Description of the changes
- added badger e2e integration test

## How was this change tested?
- tested locally 

## Checklist
- [x] I have read
https://github.com/jaegertracing/jaeger/blob/master/CONTRIBUTING_GUIDELINES.md
- [x] I have signed all commits
- [x] I have added unit tests for the new functionality
- [x] I have run lint and test steps successfully
  - for `jaeger`: `make lint test`
  - for `jaeger-ui`: `yarn lint` and `yarn test`

---------

Signed-off-by: Harshvir Potpose <hpotpose62@gmail.com>
Signed-off-by: Harshvir Potpose <122517264+akagami-harsh@users.noreply.github.com>
Signed-off-by: Yuri Shkuro <github@ysh.us>
Co-authored-by: Yuri Shkuro <yurishkuro@users.noreply.github.com>
Co-authored-by: Yuri Shkuro <github@ysh.us>

.github/workflows/ci-badger.yaml

@@ -18,6 +18,9 @@ permissions:  # added using https://github.com/step-security/secure-workflows
 jobs:
   badger:
     runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        version: [v1, v2]
     steps:
     - name: Harden Runner
       uses: step-security/harden-runner@63c24ba6bd7ba022e95695ff85de572c04a18142 # v2.7.0
@@ -31,7 +34,16 @@ jobs:
         go-version: 1.22.x
 
     - name: Run Badger storage integration tests
-      run: make badger-storage-integration-test
+      run: |
+        case ${{ matrix.version }} in
+          v1)
+              make badger-storage-integration-test
+            ;;
+          v2)
+            STORAGE=badger \
+              make jaeger-v2-storage-integration-test
+            ;;
+        esac
 
     - name: Setup CODECOV_TOKEN
       uses: ./.github/actions/setup-codecov

cmd/jaeger/badger_config.yaml

@@ -20,12 +20,14 @@ extensions:
         ephemeral: false
         maintenance_interval: 5
         metrics_update_interval: 10
+        span_store_ttl: 72h
       badger_archive:
         directory_key: "/tmp/jaeger_archive/"
         directory_value: "/tmp/jaeger_archive/"
         ephemeral: false
         maintenance_interval: 5
         metrics_update_interval: 10
+        span_store_ttl: 720h
 
 receivers:
   otlp:

cmd/jaeger/internal/components.go

@@ -29,6 +29,7 @@ import (
 	"github.com/jaegertracing/jaeger/cmd/jaeger/internal/exporters/storageexporter"
 	"github.com/jaegertracing/jaeger/cmd/jaeger/internal/extension/jaegerquery"
 	"github.com/jaegertracing/jaeger/cmd/jaeger/internal/extension/jaegerstorage"
+	"github.com/jaegertracing/jaeger/cmd/jaeger/internal/integration/storagecleaner"
 )
 
 type builders struct {
@@ -60,6 +61,7 @@ func (b builders) build() (otelcol.Factories, error) {
 		// add-ons
 		jaegerquery.NewFactory(),
 		jaegerstorage.NewFactory(),
+		storagecleaner.NewFactory(),
 		// TODO add adaptive sampling
 	)
 	if err != nil {

cmd/jaeger/internal/integration/badger_test.go

@@ -0,0 +1,51 @@
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package integration
+
+import (
+	"fmt"
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/jaegertracing/jaeger/cmd/jaeger/internal/integration/storagecleaner"
+	"github.com/jaegertracing/jaeger/plugin/storage/integration"
+)
+
+func cleanUp(t *testing.T) {
+	Addr := fmt.Sprintf("http://0.0.0.0:%s%s", storagecleaner.Port, storagecleaner.URL)
+	r, err := http.NewRequest(http.MethodPost, Addr, nil)
+	require.NoError(t, err)
+
+	client := &http.Client{}
+
+	resp, err := client.Do(r)
+	require.NoError(t, err)
+	defer resp.Body.Close()
+
+	require.Equal(t, http.StatusOK, resp.StatusCode)
+}
+
+func TestBadgerStorage(t *testing.T) {
+	integration.SkipUnlessEnv(t, "badger")
+
+	s := &E2EStorageIntegration{
+		ConfigFile: "../../badger_config.yaml",
+		StorageIntegration: integration.StorageIntegration{
+			SkipBinaryAttrs: true,
+			SkipArchiveTest: true,
+			CleanUp:         cleanUp,
+
+			// TODO: remove this once badger supports returning spanKind from GetOperations
+			// Cf https://github.com/jaegertracing/jaeger/issues/1922
+			GetOperationsMissingSpanKind: true,
+		},
+	}
+	s.e2eInitialize(t)
+	t.Cleanup(func() {
+		s.e2eCleanUp(t)
+	})
+	s.RunAll(t)
+}

cmd/jaeger/internal/integration/grpc_test.go

@@ -33,7 +33,7 @@ func TestGRPCStorage(t *testing.T) {
 	integration.SkipUnlessEnv(t, "grpc")
 
 	s := &GRPCStorageIntegration{}
-	s.ConfigFile = "cmd/jaeger/grpc_config.yaml"
+	s.ConfigFile = "../../grpc_config.yaml"
 	s.SkipBinaryAttrs = true
 
 	s.initialize(t)

cmd/jaeger/internal/integration/integration.go

@@ -7,9 +7,11 @@ import (
 	"io"
 	"os"
 	"os/exec"
+	"path/filepath"
 	"testing"
 
 	"github.com/stretchr/testify/require"
+	"gopkg.in/yaml.v3"
 
 	"github.com/jaegertracing/jaeger/pkg/testutils"
 	"github.com/jaegertracing/jaeger/plugin/storage/integration"
@@ -38,10 +40,11 @@ type E2EStorageIntegration struct {
 // This function should be called before any of the tests start.
 func (s *E2EStorageIntegration) e2eInitialize(t *testing.T) {
 	logger, _ := testutils.NewLogger()
+	configFile := createStorageCleanerConfig(t, s.ConfigFile)
 
 	cmd := exec.Cmd{
 		Path: "./cmd/jaeger/jaeger",
-		Args: []string{"jaeger", "--config", s.ConfigFile},
+		Args: []string{"jaeger", "--config", configFile},
 		// Change the working directory to the root of this project
 		// since the binary config file jaeger_query's ui_config points to
 		// "./cmd/jaeger/config-ui.json"
@@ -67,3 +70,30 @@ func (s *E2EStorageIntegration) e2eCleanUp(t *testing.T) {
 	require.NoError(t, s.SpanReader.(io.Closer).Close())
 	require.NoError(t, s.SpanWriter.(io.Closer).Close())
 }
+
+func createStorageCleanerConfig(t *testing.T, configFile string) string {
+	data, err := os.ReadFile(configFile)
+	require.NoError(t, err)
+	var config map[string]interface{}
+	err = yaml.Unmarshal(data, &config)
+	require.NoError(t, err)
+
+	service, ok := config["service"].(map[string]interface{})
+	require.True(t, ok)
+	service["extensions"] = append(service["extensions"].([]interface{}), "storage_cleaner")
+
+	extensions, ok := config["extensions"].(map[string]interface{})
+	require.True(t, ok)
+	query, ok := extensions["jaeger_query"].(map[string]interface{})
+	require.True(t, ok)
+	trace_storage := query["trace_storage"].(string)
+	extensions["storage_cleaner"] = map[string]string{"trace_storage": trace_storage}
+
+	newData, err := yaml.Marshal(config)
+	require.NoError(t, err)
+	tempFile := filepath.Join(t.TempDir(), "storageCleaner_config.yaml")
+	err = os.WriteFile(tempFile, newData, 0o600)
+	require.NoError(t, err)
+
+	return tempFile
+}

cmd/jaeger/internal/integration/receivers/storagereceiver/testdata/config.yaml

@@ -3,4 +3,4 @@ jaeger_storage_receiver/defaults:
   trace_storage: storage
 jaeger_storage_receiver/filled:
   trace_storage: storage
-  pull_interval: 2s
+  pull_interval: 2s

cmd/jaeger/internal/integration/storagecleaner/README.md

@@ -0,0 +1,47 @@
+# storage_cleaner
+
+This module implements an extension that allows purging the backend storage by making an HTTP POST request to it. 
+
+The storage_cleaner extension is intended to be used only in tests, providing a way to clear the storage between test runs. Making a POST request to the exposed endpoint will delete all data in storage.
+
+
+```mermaid
+flowchart LR
+    Receiver --> Processor
+    Processor --> Exporter
+    JaegerStorageExension -->|"(1) get storage"| Exporter
+    Exporter -->|"(2) write trace"| Storage
+
+    E2E_test -->|"(1) POST /purge"| HTTP_endpoint
+    JaegerStorageExension -->|"(2) getStorage()"| HTTP_endpoint
+    HTTP_endpoint -.->|"(3) storage.(*storage.Purger).Purge()"| Storage
+
+    subgraph Jaeger Collector
+        Receiver
+        Processor
+        Exporter
+        
+        Storage
+        StorageCleanerExtension
+        HTTP_endpoint
+        subgraph JaegerStorageExension
+            Storage
+        end
+        subgraph StorageCleanerExtension
+            HTTP_endpoint
+        end
+    end
+```
+
+# Getting Started
+
+The following settings are required:
+
+- `trace_storage` : name of a storage backend defined in `jaegerstorage` extension
+
+```yaml
+extensions:
+  storage_cleaner:
+    trace_storage: storage_name
+```
+

cmd/jaeger/internal/integration/storagecleaner/config.go

@@ -0,0 +1,18 @@
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package storagecleaner
+
+import (
+	"github.com/asaskevich/govalidator"
+)
+
+type Config struct {
+	TraceStorage string `valid:"required" mapstructure:"trace_storage"`
+	Port         string `mapstructure:"port"`
+}
+
+func (cfg *Config) Validate() error {
+	_, err := govalidator.ValidateStruct(cfg)
+	return err
+}

cmd/jaeger/internal/integration/storagecleaner/config_test.go

@@ -0,0 +1,23 @@
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package storagecleaner
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestStorageExtensionConfig(t *testing.T) {
+	config := createDefaultConfig().(*Config)
+	config.TraceStorage = "storage"
+	err := config.Validate()
+	require.NoError(t, err)
+}
+
+func TestStorageExtensionConfigError(t *testing.T) {
+	config := createDefaultConfig().(*Config)
+	err := config.Validate()
+	require.ErrorContains(t, err, "non zero value required")
+}

cmd/jaeger/internal/integration/storagecleaner/extension.go

@@ -0,0 +1,98 @@
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package storagecleaner
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/http"
+	"time"
+
+	"github.com/gorilla/mux"
+	"go.opentelemetry.io/collector/component"
+	"go.opentelemetry.io/collector/extension"
+
+	"github.com/jaegertracing/jaeger/cmd/jaeger/internal/extension/jaegerstorage"
+	"github.com/jaegertracing/jaeger/storage"
+)
+
+var (
+	_ extension.Extension = (*storageCleaner)(nil)
+	_ extension.Dependent = (*storageCleaner)(nil)
+)
+
+const (
+	Port = "9231"
+	URL  = "/purge"
+)
+
+type storageCleaner struct {
+	config   *Config
+	server   *http.Server
+	settings component.TelemetrySettings
+}
+
+func newStorageCleaner(config *Config, telemetrySettings component.TelemetrySettings) *storageCleaner {
+	return &storageCleaner{
+		config:   config,
+		settings: telemetrySettings,
+	}
+}
+
+func (c *storageCleaner) Start(ctx context.Context, host component.Host) error {
+	storageFactory, err := jaegerstorage.GetStorageFactory(c.config.TraceStorage, host)
+	if err != nil {
+		return fmt.Errorf("cannot find storage factory '%s': %w", c.config.TraceStorage, err)
+	}
+
+	purgeStorage := func() error {
+		purger, ok := storageFactory.(storage.Purger)
+		if !ok {
+			return fmt.Errorf("storage %s does not implement Purger interface", c.config.TraceStorage)
+		}
+		if err := purger.Purge(); err != nil {
+			return fmt.Errorf("error purging storage: %w", err)
+		}
+		return nil
+	}
+
+	purgeHandler := func(w http.ResponseWriter, r *http.Request) {
+		if err := purgeStorage(); err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
+		w.WriteHeader(http.StatusOK)
+		w.Write([]byte("Purge request processed successfully"))
+	}
+
+	r := mux.NewRouter()
+	r.HandleFunc(URL, purgeHandler).Methods(http.MethodPost)
+	c.server = &http.Server{
+		Addr:              ":" + c.config.Port,
+		Handler:           r,
+		ReadHeaderTimeout: 3 * time.Second,
+	}
+	go func() {
+		if err := c.server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
+			err = fmt.Errorf("error starting cleaner server: %w", err)
+			c.settings.ReportStatus(component.NewFatalErrorEvent(err))
+		}
+	}()
+
+	return nil
+}
+
+func (c *storageCleaner) Shutdown(ctx context.Context) error {
+	if c.server != nil {
+		if err := c.server.Shutdown(ctx); err != nil {
+			return fmt.Errorf("error shutting down cleaner server: %w", err)
+		}
+	}
+	return nil
+}
+
+func (c *storageCleaner) Dependencies() []component.ID {
+	return []component.ID{jaegerstorage.ID}
+}