feat: close jaeger connection (#3839)

* feat: reuse jaeger connection * invalidate cache if connection test fails on tracing step * remove cache and close datastore after poll * clean up * check for error before closing conn --------- Co-authored-by: Oscar Reyes <[email protected]>
kubeshop · May 3, 2024 · 9c3bbb0 · 9c3bbb0
1 parent 4a85d9d
commit 9c3bbb0
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 3 deletions.
diff --git a/agent/tracedb/connection/options.go b/agent/tracedb/connection/options.go
@@ -24,7 +24,7 @@ func WithAuthenticationTest(step TestStep) TesterOption {
 	}
 }
 
-func WithPollingTest(step TestStep) TesterOption {
+func WithPollingTest(step CloseableTestStep) TesterOption {
 	return func(t *Tester) {
 		t.pollingTestStep = step
 	}

diff --git a/agent/tracedb/connection/tester.go b/agent/tracedb/connection/tester.go
@@ -10,13 +10,18 @@ type TestStep interface {
 	TestConnection(ctx context.Context) model.ConnectionTestStep
 }
 
+type CloseableTestStep interface {
+	TestStep
+	CloseConnection() error
+}
+
 type TesterOption func(*Tester)
 
 type Tester struct {
 	portLinterStep         TestStep
 	connectivityTestStep   TestStep
 	authenticationTestStep TestStep
-	pollingTestStep        TestStep
+	pollingTestStep        CloseableTestStep
 }
 
 func NewTester(opts ...TesterOption) Tester {
@@ -53,6 +58,7 @@ func (t *Tester) TestConnection(ctx context.Context) (res model.ConnectionResult
 	res.FetchTraces = t.pollingTestStep.TestConnection(ctx)
 	if res.FetchTraces.Error != nil {
 		res.FetchTraces.Status = model.StatusFailed
+		t.pollingTestStep.CloseConnection()
 	}
 
 	return

diff --git a/agent/tracedb/connection/trace_polling_step.go b/agent/tracedb/connection/trace_polling_step.go
@@ -12,6 +12,7 @@ import (
 type DataStore interface {
 	GetTraceByID(context.Context, string) (traces.Trace, error)
 	GetTraceID() trace.TraceID
+	Close() error
 }
 
 type tracePollingTestStep struct {
@@ -35,6 +36,10 @@ func (s *tracePollingTestStep) TestConnection(ctx context.Context) model.Connect
 	}
 }
 
-func TracePollingTestStep(ds DataStore) TestStep {
+func (s *tracePollingTestStep) CloseConnection() error {
+	return s.dataStore.Close()
+}
+
+func TracePollingTestStep(ds DataStore) CloseableTestStep {
 	return &tracePollingTestStep{ds}
 }
diff --git a/agent/workers/poller.go b/agent/workers/poller.go
@@ -177,6 +177,7 @@ func (w *PollerWorker) poll(ctx context.Context, request *proto.PollingRequest)
 		log.Printf("Invalid datastore: %s", err.Error())
 		return err
 	}
+	defer ds.Close()
 	w.logger.Debug("Created datastore", zap.Any("datastore", ds), zap.Bool("isOTLPBasedProvider", datastoreConfig.IsOTLPBasedProvider()))
 
 	if datastoreConfig.IsOTLPBasedProvider() && w.inmemoryDatastore != nil {