From 4b2e6ca627388f5561cdf04cf9ac90e486f3eb64 Mon Sep 17 00:00:00 2001
From: Mahad Zaryab <mahadzaryab1@gmail.com>
Date: Tue, 31 Dec 2024 18:03:05 -0500
Subject: [PATCH] Fix Archive Initialization

Signed-off-by: Mahad Zaryab <mahadzaryab1@gmail.com>
---
 .../internal/extension/jaegerquery/server.go  | 28 +++++++++++------
 cmd/query/app/flags.go                        | 31 ++++++++++++-------
 2 files changed, 38 insertions(+), 21 deletions(-)

diff --git a/cmd/jaeger/internal/extension/jaegerquery/server.go b/cmd/jaeger/internal/extension/jaegerquery/server.go
index ccccaa25707..c3ccefb2a1f 100644
--- a/cmd/jaeger/internal/extension/jaegerquery/server.go
+++ b/cmd/jaeger/internal/extension/jaegerquery/server.go
@@ -164,33 +164,41 @@ func (s *server) addV2ArchiveStorage(opts *v2querysvc.QueryServiceOptions, host
 		return fmt.Errorf("cannot find archive storage factory: %w", err)
 	}
 
-	archiveFactory, ok := f.(storage.ArchiveFactory)
+	if !s.initializeV2ArchiveStorage(f, opts) {
+		s.telset.Logger.Info("Archive storage not initialized")
+	}
+	return nil
+}
+
+func (s *server) initializeV2ArchiveStorage(
+	storageFactory storage.Factory,
+	opts *v2querysvc.QueryServiceOptions) bool {
+	archiveFactory, ok := storageFactory.(storage.ArchiveFactory)
 	if !ok {
 		s.telset.Logger.Info("Archive storage not supported by the factory")
+		return false
 	}
 	reader, err := archiveFactory.CreateArchiveSpanReader()
 	if errors.Is(err, storage.ErrArchiveStorageNotConfigured) || errors.Is(err, storage.ErrArchiveStorageNotSupported) {
 		s.telset.Logger.Info("Archive storage not created", zap.String("reason", err.Error()))
+		return false
 	}
 	if err != nil {
 		s.telset.Logger.Error("Cannot init archive storage reader", zap.Error(err))
+		return false
 	}
 	writer, err := archiveFactory.CreateArchiveSpanWriter()
 	if errors.Is(err, storage.ErrArchiveStorageNotConfigured) || errors.Is(err, storage.ErrArchiveStorageNotSupported) {
 		s.telset.Logger.Info("Archive storage not created", zap.String("reason", err.Error()))
+		return false
 	}
 	if err != nil {
 		s.telset.Logger.Error("Cannot init archive storage writer", zap.Error(err))
+		return false
 	}
-
-	if writer != nil && reader != nil {
-		opts.ArchiveTraceReader = v1adapter.NewTraceReader(reader)
-		opts.ArchiveTraceWriter = v1adapter.NewTraceWriter(writer)
-	} else {
-		s.telset.Logger.Info("Archive storage not initialized")
-	}
-
-	return nil
+	opts.ArchiveTraceReader = v1adapter.NewTraceReader(reader)
+	opts.ArchiveTraceWriter = v1adapter.NewTraceWriter(writer)
+	return true
 }
 
 func (s *server) createMetricReader(host component.Host) (metricstore.Reader, error) {
diff --git a/cmd/query/app/flags.go b/cmd/query/app/flags.go
index 74543161ed4..134637e57c0 100644
--- a/cmd/query/app/flags.go
+++ b/cmd/query/app/flags.go
@@ -152,36 +152,45 @@ func (qOpts *QueryOptions) BuildQueryServiceOptions(storageFactory storage.BaseF
 
 func (qOpts *QueryOptions) BuildV2QueryServiceOptions(storageFactory storage.BaseFactory, logger *zap.Logger) *v2querysvc.QueryServiceOptions {
 	opts := &v2querysvc.QueryServiceOptions{}
+	if !initializeV2ArchiveStorage(storageFactory, opts, logger) {
+		logger.Info("Archive storage not initialized")
+	}
+
+	opts.Adjuster = v2adjuster.Sequence(v2adjuster.StandardAdjusters(qOpts.MaxClockSkewAdjust)...)
 
+	return opts
+}
+
+func initializeV2ArchiveStorage(
+	storageFactory storage.BaseFactory,
+	opts *v2querysvc.QueryServiceOptions,
+	logger *zap.Logger) bool {
 	archiveFactory, ok := storageFactory.(storage.ArchiveFactory)
 	if !ok {
 		logger.Info("Archive storage not supported by the factory")
+		return false
 	}
 	reader, err := archiveFactory.CreateArchiveSpanReader()
 	if errors.Is(err, storage.ErrArchiveStorageNotConfigured) || errors.Is(err, storage.ErrArchiveStorageNotSupported) {
 		logger.Info("Archive storage not created", zap.String("reason", err.Error()))
+		return false
 	}
 	if err != nil {
 		logger.Error("Cannot init archive storage reader", zap.Error(err))
+		return false
 	}
 	writer, err := archiveFactory.CreateArchiveSpanWriter()
 	if errors.Is(err, storage.ErrArchiveStorageNotConfigured) || errors.Is(err, storage.ErrArchiveStorageNotSupported) {
 		logger.Info("Archive storage not created", zap.String("reason", err.Error()))
+		return false
 	}
 	if err != nil {
 		logger.Error("Cannot init archive storage writer", zap.Error(err))
+		return false
 	}
-
-	if writer != nil && reader != nil {
-		opts.ArchiveTraceReader = v1adapter.NewTraceReader(reader)
-		opts.ArchiveTraceWriter = v1adapter.NewTraceWriter(writer)
-	} else {
-		logger.Info("Archive storage not initialized")
-	}
-
-	opts.Adjuster = v2adjuster.Sequence(v2adjuster.StandardAdjusters(qOpts.MaxClockSkewAdjust)...)
-
-	return opts
+	opts.ArchiveTraceReader = v1adapter.NewTraceReader(reader)
+	opts.ArchiveTraceWriter = v1adapter.NewTraceWriter(writer)
+	return true
 }
 
 // stringSliceAsHeader parses a slice of strings and returns a http.Header.