From 55089ae56750e93893f584bcb5486dceb9e37812 Mon Sep 17 00:00:00 2001
From: Will Sewell <willsewell@monzo.com>
Date: Sun, 31 Mar 2024 22:41:11 +0100
Subject: [PATCH] Fix goroutine leaks in integration/cassandra_test.go

Signed-off-by: Will Sewell <willsewell@monzo.com>
---
 plugin/storage/cassandra/factory.go          |  7 +++++++
 plugin/storage/integration/cassandra_test.go | 19 +++++++++++--------
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/plugin/storage/cassandra/factory.go b/plugin/storage/cassandra/factory.go
index dfa12fcb974..664d5a17259 100644
--- a/plugin/storage/cassandra/factory.go
+++ b/plugin/storage/cassandra/factory.go
@@ -232,6 +232,13 @@ var _ io.Closer = (*Factory)(nil)
 
 // Close closes the resources held by the factory
 func (f *Factory) Close() error {
+	if f.primarySession != nil {
+		f.primarySession.Close()
+	}
+	if f.archiveSession != nil {
+		f.archiveSession.Close()
+	}
+
 	f.Options.Get(archiveStorageConfig)
 	if cfg := f.Options.Get(archiveStorageConfig); cfg != nil {
 		cfg.TLS.Close()
diff --git a/plugin/storage/integration/cassandra_test.go b/plugin/storage/integration/cassandra_test.go
index facbb11a5f4..ff7ee0e04c3 100644
--- a/plugin/storage/integration/cassandra_test.go
+++ b/plugin/storage/integration/cassandra_test.go
@@ -18,6 +18,7 @@ package integration
 import (
 	"errors"
 	"fmt"
+	"io"
 	"os"
 	"testing"
 
@@ -81,28 +82,28 @@ func (s *CassandraStorageIntegration) initializeCassandraFactory(flags []string)
 	return f, nil
 }
 
-func (s *CassandraStorageIntegration) initializeCassandra() error {
+func (s *CassandraStorageIntegration) initializeCassandra() (io.Closer, error) {
 	f, err := s.initializeCassandraFactory([]string{
 		"--cassandra.keyspace=jaeger_v1_dc1",
 	})
 	if err != nil {
-		return err
+		return nil, err
 	}
 	s.session = f.PrimarySession()
 	if s.SpanWriter, err = f.CreateSpanWriter(); err != nil {
-		return err
+		return nil, err
 	}
 	if s.SpanReader, err = f.CreateSpanReader(); err != nil {
-		return err
+		return nil, err
 	}
 	if s.SamplingStore, err = f.CreateSamplingStore(0); err != nil {
-		return err
+		return nil, err
 	}
 
 	if err = s.initializeDependencyReaderAndWriter(f); err != nil {
-		return err
+		return nil, err
 	}
-	return nil
+	return f, nil
 }
 
 func (s *CassandraStorageIntegration) initializeDependencyReaderAndWriter(f *cassandra.Factory) error {
@@ -125,6 +126,8 @@ func TestCassandraStorage(t *testing.T) {
 		t.Skip("Integration test against Cassandra skipped; set STORAGE env var to cassandra to run this")
 	}
 	s := newCassandraStorageIntegration()
-	require.NoError(t, s.initializeCassandra())
+	closer, err := s.initializeCassandra()
+	require.NoError(t, err)
 	s.IntegrationTestAll(t)
+	require.NoError(t, closer.Close())
 }