From 8a283a6f5e04641f0b314f5e90b0429b1829981e Mon Sep 17 00:00:00 2001
From: Zeger-Jan van de Weg <git@zjvandeweg.nl>
Date: Thu, 7 Mar 2019 12:02:23 +0100
Subject: [PATCH] Allow custom logger for the tracing pacakge

The logger defaulted to the go standard library logger, but that doesn't
allow for structured logging. When booting multiple LabKit applications
with tracing that would mean that multiple logging formats were used for
applications which broke my poor mans log filter (jq reading from a
stream).
---
 tracing/env_injector.go                |  5 ++---
 tracing/env_injector_option.go         | 17 +++++++++++++++--
 tracing/initialization.go              |  3 ++-
 tracing/initialization_options.go      | 16 ++++++++++++++++
 tracing/initialization_options_test.go | 21 +++++++++++++++++++++
 5 files changed, 56 insertions(+), 6 deletions(-)
 create mode 100644 tracing/initialization_options_test.go

diff --git a/tracing/env_injector.go b/tracing/env_injector.go
index 30d39a1b..5ce12af6 100644
--- a/tracing/env_injector.go
+++ b/tracing/env_injector.go
@@ -3,7 +3,6 @@ package tracing
 import (
 	"context"
 	"fmt"
-	"log"
 	"os"
 	"sort"
 
@@ -22,7 +21,7 @@ type EnvInjector func(ctx context.Context, env []string) []string
 
 // NewEnvInjector will create a new environment injector
 func NewEnvInjector(opts ...EnvInjectorOption) EnvInjector {
-	/* config not yet used */ applyEnvInjectorOptions(opts)
+	conf := applyEnvInjectorOptions(opts)
 
 	return func(ctx context.Context, env []string) []string {
 		envMap := map[string]string{}
@@ -50,7 +49,7 @@ func NewEnvInjector(opts ...EnvInjectorOption) EnvInjector {
 		err := span.Tracer().Inject(span.Context(), opentracing.TextMap, carrier)
 
 		if err != nil {
-			log.Printf("tracing span injection failed: %v", err)
+			conf.log.Printf("tracing span injection failed: %v", err)
 		}
 
 		return appendMapToEnv(env, envMap)
diff --git a/tracing/env_injector_option.go b/tracing/env_injector_option.go
index e3298a4e..ec56dcab 100644
--- a/tracing/env_injector_option.go
+++ b/tracing/env_injector_option.go
@@ -1,15 +1,28 @@
 package tracing
 
-type envInjectorConfig struct{}
+import (
+	"log"
+	"os"
+)
+
+type envInjectorConfig struct {
+	log Logger
+}
 
 // EnvInjectorOption will configure an environment injector
 type EnvInjectorOption func(*envInjectorConfig)
 
 func applyEnvInjectorOptions(opts []EnvInjectorOption) envInjectorConfig {
-	config := envInjectorConfig{}
+	config := envInjectorConfig{log: log.New(os.Stdout, "", log.LstdFlags)}
 	for _, v := range opts {
 		v(&config)
 	}
 
 	return config
 }
+
+func WithEnvLogger(logger Logger) EnvInjectorOption {
+	return func(config *envInjectorConfig) {
+		config.log = logger
+	}
+}
diff --git a/tracing/initialization.go b/tracing/initialization.go
index d0ec3ba0..43ec803b 100644
--- a/tracing/initialization.go
+++ b/tracing/initialization.go
@@ -2,7 +2,6 @@ package tracing
 
 import (
 	"io"
-	"log"
 
 	opentracing "github.com/opentracing/opentracing-go"
 	"gitlab.com/gitlab-org/labkit/tracing/connstr"
@@ -23,6 +22,8 @@ func Initialize(opts ...InitializationOption) io.Closer {
 		return &nopCloser{}
 	}
 
+	log := config.log
+
 	driverName, options, err := connstr.Parse(config.connectionString)
 	if err != nil {
 		log.Printf("unable to parse connection: %v", err)
diff --git a/tracing/initialization_options.go b/tracing/initialization_options.go
index cfb4996d..187b051d 100644
--- a/tracing/initialization_options.go
+++ b/tracing/initialization_options.go
@@ -1,6 +1,7 @@
 package tracing
 
 import (
+	"log"
 	"os"
 	"path"
 )
@@ -11,6 +12,7 @@ const tracingEnvKey = "GITLAB_TRACING"
 type initializationConfig struct {
 	serviceName      string
 	connectionString string
+	log              Logger
 }
 
 // InitializationOption will configure a correlation handler
@@ -20,6 +22,7 @@ func applyInitializationOptions(opts []InitializationOption) initializationConfi
 	config := initializationConfig{
 		serviceName:      path.Base(os.Args[0]),
 		connectionString: os.Getenv(tracingEnvKey),
+		log:              log.New(os.Stdout, "", log.LstdFlags),
 	}
 	for _, v := range opts {
 		v(&config)
@@ -43,3 +46,16 @@ func WithConnectionString(connectionString string) InitializationOption {
 		config.connectionString = connectionString
 	}
 }
+
+// The logger interface allows for customization of the logger with formatting
+type Logger interface {
+	Printf(format string, args ...interface{})
+}
+
+// WithInitLogger overrides the default logger and uses the any logger that implements
+// the Logger interface of this package
+func WithInitLogger(logger Logger) InitializationOption {
+	return func(config *initializationConfig) {
+		config.log = logger
+	}
+}
diff --git a/tracing/initialization_options_test.go b/tracing/initialization_options_test.go
new file mode 100644
index 00000000..fa58b165
--- /dev/null
+++ b/tracing/initialization_options_test.go
@@ -0,0 +1,21 @@
+package tracing
+
+import (
+	"bytes"
+	"log"
+	"testing"
+)
+
+func TestCustomLogger(t *testing.T) {
+	var buf bytes.Buffer
+	cLog := log.New(&buf, "", 0)
+
+	conf := applyInitializationOptions([]InitializationOption{WithLogger(cLog)})
+
+	message := "mepmep"
+	conf.log.Printf(message)
+
+	if buf.String() != message+"\n" {
+		t.Fatalf("custom logger did not log %s", message)
+	}
+}
-- 
GitLab