invalidate cache if connection test fails on tracing step

kubeshop · May 2, 2024 · 44e081f · 44e081f
1 parent e3d1f65
commit 44e081f
Show file tree

Hide file tree

Showing 4 changed files with 22 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/tracedb/jaegerdb.go b/agent/tracedb/jaegerdb.go
@@ -43,6 +43,13 @@ func saveConnectionToCache(grpcConfig *datastore.GRPCClientSettings, traceDB Tra
 	connectionPool[grpcConfig.Endpoint] = traceDB
 }
 
+func invalidateConnectionCache(endpoint string) {
+	connectionPoolMutex.Lock()
+	defer connectionPoolMutex.Unlock()
+
+	delete(connectionPool, endpoint)
+}
+
 type jaegerTraceDB struct {
 	realTraceDB
 	dataSource datasource.DataSource
@@ -107,6 +114,7 @@ func (jtd *jaegerTraceDB) Ready() bool {
 }
 
 func (jtd *jaegerTraceDB) Close() error {
+	invalidateConnectionCache(jtd.dataSource.Endpoint())
 	return jtd.dataSource.Close()
 }