From 9c3bbb0955ab8f531d08d56e28a3389e562c8d03 Mon Sep 17 00:00:00 2001
From: Matheus Nogueira <matheus.nogueira2008@gmail.com>
Date: Fri, 3 May 2024 13:23:23 -0300
Subject: [PATCH] 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 <oscar-rreyes1@hotmail.com>
---
 agent/tracedb/connection/options.go            | 2 +-
 agent/tracedb/connection/tester.go             | 8 +++++++-
 agent/tracedb/connection/trace_polling_step.go | 7 ++++++-
 agent/workers/poller.go                        | 1 +
 4 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/agent/tracedb/connection/options.go b/agent/tracedb/connection/options.go
index 136cbbaef2..ceb9161deb 100644
--- 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
index 80a102fe35..6593bcb215 100644
--- 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
index 36751cc45f..987493676f 100644
--- 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
index e70bed914b..7c8b18aa5c 100644
--- 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 {