From 61e97f0896a2e23ebdd3e2623bdbd3d1a0c6f4cc Mon Sep 17 00:00:00 2001
From: Gabriel Adrian Samfira <gsamfira@cloudbasesolutions.com>
Date: Sat, 6 Jan 2024 00:21:50 +0000
Subject: [PATCH] Append pool_type and pool_mgr info to logs

Pool managers will have 2 fields identifying which manager generated
the log line.

In the future, we will add tracking ids in various cases, allowing
us to track down issues faster.

Signed-off-by: Gabriel Adrian Samfira <gsamfira@cloudbasesolutions.com>
---
 cmd/garm/main.go            |  9 ++++++---
 runner/pool/enterprise.go   |  1 +
 runner/pool/organization.go |  1 +
 runner/pool/repository.go   |  1 +
 runner/pool/util.go         |  9 ---------
 util/logging.go             | 30 ++++++++++++++++++++++++++++++
 6 files changed, 39 insertions(+), 12 deletions(-)
 create mode 100644 util/logging.go

diff --git a/cmd/garm/main.go b/cmd/garm/main.go
index 008adf3c..0d325750 100644
--- a/cmd/garm/main.go
+++ b/cmd/garm/main.go
@@ -37,6 +37,7 @@ import (
 	"github.com/cloudbase/garm/database/common"
 	"github.com/cloudbase/garm/metrics"
 	"github.com/cloudbase/garm/runner"
+	garmUtil "github.com/cloudbase/garm/util"
 	"github.com/cloudbase/garm/util/appdefaults"
 	"github.com/cloudbase/garm/websocket"
 	lumberjack "gopkg.in/natefinch/lumberjack.v2"
@@ -105,8 +106,6 @@ func setupLogging(ctx context.Context, cfg *config.Config, hub *websocket.Hub) {
 	}
 
 	wr := io.MultiWriter(writers...)
-	// TODO: delete this once we migrate to slog
-	log.SetOutput(wr)
 
 	logCfg := cfg.GetLoggingConfig()
 	var logLevel slog.Level
@@ -136,7 +135,11 @@ func setupLogging(ctx context.Context, cfg *config.Config, hub *websocket.Hub) {
 	default:
 		han = slog.NewTextHandler(wr, &opts)
 	}
-	slog.SetDefault(slog.New(han))
+
+	wrapped := garmUtil.ContextHandler{
+		Handler: han,
+	}
+	slog.SetDefault(slog.New(wrapped))
 
 }
 
diff --git a/runner/pool/enterprise.go b/runner/pool/enterprise.go
index fa0e13be..f03947b5 100644
--- a/runner/pool/enterprise.go
+++ b/runner/pool/enterprise.go
@@ -25,6 +25,7 @@ import (
 var _ poolHelper = &enterprise{}
 
 func NewEnterprisePoolManager(ctx context.Context, cfg params.Enterprise, cfgInternal params.Internal, providers map[string]common.Provider, store dbCommon.Store) (common.PoolManager, error) {
+	ctx = util.WithContext(ctx, slog.Any("pool_mgr", cfg.Name), slog.Any("pool_type", params.EnterprisePool))
 	ghc, ghEnterpriseClient, err := util.GithubClient(ctx, cfgInternal.OAuth2Token, cfgInternal.GithubCredentialsDetails)
 	if err != nil {
 		return nil, errors.Wrap(err, "getting github client")
diff --git a/runner/pool/organization.go b/runner/pool/organization.go
index 88aa1632..0f534508 100644
--- a/runner/pool/organization.go
+++ b/runner/pool/organization.go
@@ -39,6 +39,7 @@ import (
 var _ poolHelper = &organization{}
 
 func NewOrganizationPoolManager(ctx context.Context, cfg params.Organization, cfgInternal params.Internal, providers map[string]common.Provider, store dbCommon.Store) (common.PoolManager, error) {
+	ctx = util.WithContext(ctx, slog.Any("pool_mgr", cfg.Name), slog.Any("pool_type", params.OrganizationPool))
 	ghc, _, err := util.GithubClient(ctx, cfgInternal.OAuth2Token, cfgInternal.GithubCredentialsDetails)
 	if err != nil {
 		return nil, errors.Wrap(err, "getting github client")
diff --git a/runner/pool/repository.go b/runner/pool/repository.go
index f41d6d5f..0505eece 100644
--- a/runner/pool/repository.go
+++ b/runner/pool/repository.go
@@ -39,6 +39,7 @@ import (
 var _ poolHelper = &repository{}
 
 func NewRepositoryPoolManager(ctx context.Context, cfg params.Repository, cfgInternal params.Internal, providers map[string]common.Provider, store dbCommon.Store) (common.PoolManager, error) {
+	ctx = util.WithContext(ctx, slog.Any("pool_mgr", fmt.Sprintf("%s/%s", cfg.Owner, cfg.Name)), slog.Any("pool_type", params.RepositoryPool))
 	ghc, _, err := util.GithubClient(ctx, cfgInternal.OAuth2Token, cfgInternal.GithubCredentialsDetails)
 	if err != nil {
 		return nil, errors.Wrap(err, "getting github client")
diff --git a/runner/pool/util.go b/runner/pool/util.go
index 4a8c09e3..b97bf26f 100644
--- a/runner/pool/util.go
+++ b/runner/pool/util.go
@@ -1,7 +1,6 @@
 package pool
 
 import (
-	"log"
 	"sort"
 	"strings"
 	"sync"
@@ -57,11 +56,3 @@ func (p *poolsForTags) Add(tags []string, pools []params.Pool) *poolRoundRobin {
 	v, _ := p.pools.LoadOrStore(key, poolRR)
 	return v.(*poolRoundRobin)
 }
-
-func (r *basePoolManager) log(msg string, args ...interface{}) {
-	msgArgs := []interface{}{
-		r.helper.String(),
-	}
-	msgArgs = append(msgArgs, args...)
-	log.Printf("[Pool mgr %s] "+msg, msgArgs...)
-}
diff --git a/util/logging.go b/util/logging.go
new file mode 100644
index 00000000..ac35863b
--- /dev/null
+++ b/util/logging.go
@@ -0,0 +1,30 @@
+package util
+
+import (
+	"context"
+	"log/slog"
+)
+
+type slogContextKey string
+
+const (
+	slogCtxFields slogContextKey = "slog_ctx_fields"
+)
+
+type ContextHandler struct {
+	slog.Handler
+}
+
+func (h ContextHandler) Handle(ctx context.Context, r slog.Record) error {
+	attrs, ok := ctx.Value(slogCtxFields).([]slog.Attr)
+	if ok {
+		for _, v := range attrs {
+			r.AddAttrs(v)
+		}
+	}
+	return h.Handler.Handle(ctx, r)
+}
+
+func WithContext(ctx context.Context, attrs ...slog.Attr) context.Context {
+	return context.WithValue(ctx, slogCtxFields, attrs)
+}