From c5ac4b28cc3a9b64e4bc2da054f5c08396ee2cd6 Mon Sep 17 00:00:00 2001
From: rene <41963722+renaynay@users.noreply.github.com>
Date: Wed, 5 Oct 2022 15:22:56 +0200
Subject: [PATCH] feat: Allow annotation on subsystem names after instantiation

---
 log.go      | 19 +++++++++++++++++++
 log_test.go | 27 +++++++++++++++++++++++++++
 2 files changed, 46 insertions(+)
 create mode 100644 log_test.go

diff --git a/log.go b/log.go
index 81a7fcb..84d5506 100644
--- a/log.go
+++ b/log.go
@@ -3,6 +3,7 @@
 package log
 
 import (
+	"fmt"
 	"time"
 
 	"go.uber.org/zap"
@@ -58,6 +59,24 @@ type ZapEventLogger struct {
 	system     string
 }
 
+// AnnotateName annotates the logger's system name with the given annotation, prepending
+// it if specified.
+func (logger *ZapEventLogger) AnnotateName(annotation string, prepend bool) {
+	newName := logger.system
+	if prepend {
+		newName = fmt.Sprintf("%s:%s", annotation, newName)
+	} else {
+		newName = fmt.Sprintf("%s:%s", newName, annotation)
+	}
+	// register new logger system name
+	loggerMutex.Lock()
+	defer loggerMutex.Unlock()
+	delete(loggers, logger.system)
+	loggers[newName] = &logger.SugaredLogger
+
+	logger.system = newName
+}
+
 // Warning is for compatibility
 // Deprecated: use Warn(args ...interface{}) instead
 func (logger *ZapEventLogger) Warning(args ...interface{}) {
diff --git a/log_test.go b/log_test.go
new file mode 100644
index 0000000..44fbf71
--- /dev/null
+++ b/log_test.go
@@ -0,0 +1,27 @@
+package log
+
+import (
+	"fmt"
+	"testing"
+)
+
+// TestLogger_AnnotateName ensures that logger's system names are
+// accurately annotated.
+func TestLogger_AnnotateName(t *testing.T) {
+	// set up several loggers
+	Logger("test1")
+	Logger("test2")
+	Logger("test3")
+
+	annotation := "annotation"
+
+	for _, sub := range GetSubsystems() {
+		Logger(sub).AnnotateName(annotation, true)
+	}
+
+	for _, sub := range GetSubsystems() {
+		if sub[:len(annotation)] != annotation {
+			t.Fatalf("expected %s, got %s", fmt.Sprintf("%s:%s", annotation, sub), sub)
+		}
+	}
+}