From 347e3e37cf55244a918b536483e84dcfe15a7dc3 Mon Sep 17 00:00:00 2001
From: Dylan Guedes <djmgguedes@gmail.com>
Date: Fri, 26 Nov 2021 08:38:41 -0300
Subject: [PATCH] Loki: Implement custom /config handler (#4785)

* Implement custom /config handler.

* Make `customConfigEndpointHandlerFn` public.
---
 cmd/loki/main.go      |  2 +-
 pkg/loki/loki.go      | 21 ++++++++++--
 pkg/loki/loki_test.go | 80 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 99 insertions(+), 4 deletions(-)

diff --git a/cmd/loki/main.go b/cmd/loki/main.go
index cf7ce4768ca0..0b4a735e334e 100644
--- a/cmd/loki/main.go
+++ b/cmd/loki/main.go
@@ -93,6 +93,6 @@ func main() {
 
 	level.Info(util_log.Logger).Log("msg", "Starting Loki", "version", version.Info())
 
-	err = t.Run()
+	err = t.Run(loki.RunOpts{})
 	util_log.CheckFatal("running loki", err)
 }
diff --git a/pkg/loki/loki.go b/pkg/loki/loki.go
index 2afa07f5d787..f4f57ea2ae51 100644
--- a/pkg/loki/loki.go
+++ b/pkg/loki/loki.go
@@ -280,8 +280,23 @@ func newDefaultConfig() *Config {
 	return defaultConfig
 }
 
+// RunOpts configures custom behavior for running Loki.
+type RunOpts struct {
+	// CustomConfigEndpointHandlerFn is the handlerFunc to be used by the /config endpoint.
+	// If empty, default handlerFunc will be used.
+	CustomConfigEndpointHandlerFn func(http.ResponseWriter, *http.Request)
+}
+
+func (t *Loki) bindConfigEndpoint(opts RunOpts) {
+	configEndpointHandlerFn := configHandler(t.Cfg, newDefaultConfig())
+	if opts.CustomConfigEndpointHandlerFn != nil {
+		configEndpointHandlerFn = opts.CustomConfigEndpointHandlerFn
+	}
+	t.Server.HTTP.Path("/config").Methods("GET").HandlerFunc(configEndpointHandlerFn)
+}
+
 // Run starts Loki running, and blocks until a Loki stops.
-func (t *Loki) Run() error {
+func (t *Loki) Run(opts RunOpts) error {
 	serviceMap, err := t.ModuleManager.InitModuleServices(t.Cfg.Target...)
 	if err != nil {
 		return err
@@ -306,8 +321,8 @@ func (t *Loki) Run() error {
 
 	grpc_health_v1.RegisterHealthServer(t.Server.GRPC, grpcutil.NewHealthCheck(sm))
 
-	// This adds a way to see the config and the changes compared to the defaults
-	t.Server.HTTP.Path("/config").Methods("GET").HandlerFunc(configHandler(t.Cfg, newDefaultConfig()))
+	// Config endpoint adds a way to see the config and the changes compared to the defaults.
+	t.bindConfigEndpoint(opts)
 
 	// Each component serves its version.
 	t.Server.HTTP.Path("/loki/api/v1/status/buildinfo").Methods("GET").HandlerFunc(versionHandler())
diff --git a/pkg/loki/loki_test.go b/pkg/loki/loki_test.go
index e02fb184c484..31757ab8ed2e 100644
--- a/pkg/loki/loki_test.go
+++ b/pkg/loki/loki_test.go
@@ -3,7 +3,10 @@ package loki
 import (
 	"bytes"
 	"flag"
+	"fmt"
 	"io"
+	"io/ioutil"
+	"net/http"
 	"strings"
 	"testing"
 	"time"
@@ -86,3 +89,80 @@ func TestLoki_isModuleEnabled(t1 *testing.T) {
 		})
 	}
 }
+
+func TestLoki_CustomRunOptsBehavior(t *testing.T) {
+	yamlConfig := `target: querier
+server:
+  http_listen_port: 3100
+common:
+  path_prefix: /tmp/loki
+  ring:
+    kvstore:
+      store: inmemory
+schema_config:
+  configs:
+    - from: 2020-10-24
+      store: boltdb-shipper
+      row_shards: 10
+      object_store: filesystem
+      schema: v11
+      index:
+        prefix: index_
+        period: 24h`
+
+	cfgWrapper, _, err := configWrapperFromYAML(t, yamlConfig, nil)
+	require.NoError(t, err)
+
+	loki, err := New(cfgWrapper.Config)
+	require.NoError(t, err)
+
+	lokiHealthCheck := func() error {
+		// wait for Loki HTTP server to be ready.
+		// retries at most 10 times (1 second in total) to avoid infinite loops when no timeout is set.
+		for i := 0; i < 10; i++ {
+			// waits until request to /ready doesn't error.
+			resp, err := http.DefaultClient.Get("http://localhost:3100/ready")
+			if err != nil {
+				time.Sleep(time.Millisecond * 200)
+				continue
+			}
+
+			// waits until /ready returns OK.
+			if resp.StatusCode != http.StatusOK {
+				time.Sleep(time.Millisecond * 200)
+				continue
+			}
+
+			// Loki is healthy.
+			return nil
+		}
+
+		return fmt.Errorf("loki HTTP not healthy")
+	}
+
+	customHandlerInvoked := false
+	customHandler := func(w http.ResponseWriter, _ *http.Request) {
+		customHandlerInvoked = true
+		_, err := w.Write([]byte("abc"))
+		require.NoError(t, err)
+	}
+
+	// Run Loki querier in a different go routine and with custom /config handler.
+	go func() {
+		err := loki.Run(RunOpts{CustomConfigEndpointHandlerFn: customHandler})
+		require.NoError(t, err)
+	}()
+
+	err = lokiHealthCheck()
+	require.NoError(t, err)
+
+	resp, err := http.DefaultClient.Get("http://localhost:3100/config")
+	require.NoError(t, err)
+
+	defer resp.Body.Close()
+
+	bBytes, err := ioutil.ReadAll(resp.Body)
+	require.NoError(t, err)
+	require.Equal(t, string(bBytes), "abc")
+	assert.True(t, customHandlerInvoked)
+}