Generate and log correlation IDs

This will make it easier to tie an SSH access request to Rails API and
Gitaly requests.

3 files changed, 31 insertions, 10 deletions

diff --git a/client/client_test.go b/client/client_test.go
index 5c5fda8..d520bbb 100644
--- a/client/client_test.go
+++ b/client/client_test.go
@@ -130,6 +130,7 @@ func testSuccessfulGet(t *testing.T, client *GitlabNetClient) {
 		assert.Contains(t, entries[0].Message, "method=GET")
 		assert.Contains(t, entries[0].Message, "status=200")
 		assert.Contains(t, entries[0].Message, "Finished HTTP request")
+		assert.Contains(t, entries[0].Message, "correlation_id=")
 	})
 }
 
@@ -155,6 +156,7 @@ func testSuccessfulPost(t *testing.T, client *GitlabNetClient) {
 		assert.Contains(t, entries[0].Message, "method=POST")
 		assert.Contains(t, entries[0].Message, "status=200")
 		assert.Contains(t, entries[0].Message, "Finished HTTP request")
+		assert.Contains(t, entries[0].Message, "correlation_id=")
 	})
 }
 
@@ -172,6 +174,7 @@ func testMissing(t *testing.T, client *GitlabNetClient) {
 		assert.Contains(t, entries[0].Message, "method=GET")
 		assert.Contains(t, entries[0].Message, "status=404")
 		assert.Contains(t, entries[0].Message, "Internal API error")
+		assert.Contains(t, entries[0].Message, "correlation_id=")
 	})
 
 	t.Run("Missing error for POST", func(t *testing.T) {
@@ -187,6 +190,7 @@ func testMissing(t *testing.T, client *GitlabNetClient) {
 		assert.Contains(t, entries[0].Message, "method=POST")
 		assert.Contains(t, entries[0].Message, "status=404")
 		assert.Contains(t, entries[0].Message, "Internal API error")
+		assert.Contains(t, entries[0].Message, "correlation_id=")
 	})
 }
 
@@ -219,6 +223,7 @@ func testBrokenRequest(t *testing.T, client *GitlabNetClient) {
 		assert.Contains(t, entries[0].Message, "method=GET")
 		assert.NotContains(t, entries[0].Message, "status=")
 		assert.Contains(t, entries[0].Message, "Internal API unreachable")
+		assert.Contains(t, entries[0].Message, "correlation_id=")
 	})
 
 	t.Run("Broken request for POST", func(t *testing.T) {
@@ -235,6 +240,7 @@ func testBrokenRequest(t *testing.T, client *GitlabNetClient) {
 		assert.Contains(t, entries[0].Message, "method=POST")
 		assert.NotContains(t, entries[0].Message, "status=")
 		assert.Contains(t, entries[0].Message, "Internal API unreachable")
+		assert.Contains(t, entries[0].Message, "correlation_id=")
 	})
 }
 
diff --git a/client/gitlabnet.go b/client/gitlabnet.go
index 7fbf63e..0657ca0 100644
--- a/client/gitlabnet.go
+++ b/client/gitlabnet.go
@@ -2,6 +2,7 @@ package client
 
 import (
 	"bytes"
+	"context"
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
@@ -11,6 +12,7 @@ import (
 	"time"
 
 	log "github.com/sirupsen/logrus"
+	"gitlab.com/gitlab-org/labkit/correlation"
 )
 
 const (
@@ -57,23 +59,32 @@ func normalizePath(path string) string {
 	return path
 }
 
-func newRequest(method, host, path string, data interface{}) (*http.Request, error) {
+func newRequest(method, host, path string, data interface{}) (*http.Request, string, error) {
 	var jsonReader io.Reader
 	if data != nil {
 		jsonData, err := json.Marshal(data)
 		if err != nil {
-			return nil, err
+			return nil, "", err
 		}
 
 		jsonReader = bytes.NewReader(jsonData)
 	}
 
-	request, err := http.NewRequest(method, host+path, jsonReader)
+	correlationID, err := correlation.RandomID()
+	ctx := context.Background()
+
 	if err != nil {
-		return nil, err
+		log.WithError(err).Warn("unable to generate correlation ID")
+	} else {
+		ctx = correlation.ContextWithCorrelation(ctx, correlationID)
+	}
+
+	request, err := http.NewRequestWithContext(ctx, method, host+path, jsonReader)
+	if err != nil {
+		return nil, "", err
 	}
 
-	return request, nil
+	return request, correlationID, nil
 }
 
 func parseError(resp *http.Response) error {
@@ -100,7 +111,7 @@ func (c *GitlabNetClient) Post(path string, data interface{}) (*http.Response, e
 }
 
 func (c *GitlabNetClient) DoRequest(method, path string, data interface{}) (*http.Response, error) {
-	request, err := newRequest(method, c.httpClient.Host, path, data)
+	request, correlationID, err := newRequest(method, c.httpClient.Host, path, data)
 	if err != nil {
 		return nil, err
 	}
@@ -119,9 +130,10 @@ func (c *GitlabNetClient) DoRequest(method, path string, data interface{}) (*htt
 	start := time.Now()
 	response, err := c.httpClient.Do(request)
 	fields := log.Fields{
-		"method":      method,
-		"url":         request.URL.String(),
-		"duration_ms": time.Since(start) / time.Millisecond,
+		"correlation_id": correlationID,
+		"method":         method,
+		"url":            request.URL.String(),
+		"duration_ms":    time.Since(start) / time.Millisecond,
 	}
 	logger := log.WithFields(fields)
 
diff --git a/client/httpclient.go b/client/httpclient.go
index ff0cc25..63386f7 100644
--- a/client/httpclient.go
+++ b/client/httpclient.go
@@ -10,6 +10,8 @@ import (
 	"path/filepath"
 	"strings"
 	"time"
+
+	"gitlab.com/gitlab-org/labkit/correlation"
 )
 
 const (
@@ -39,8 +41,9 @@ func NewHTTPClient(gitlabURL, caFile, caPath string, selfSignedCert bool, readTi
 		return nil
 	}
 
+
 	c := &http.Client{
-		Transport: transport,
+		Transport: correlation.NewInstrumentedRoundTripper(transport),
 		Timeout:   readTimeout(readTimeoutSeconds),
 	}