feat: add structured log writer (#246)

funcframework/framework.go

@@ -198,7 +198,7 @@ func wrapHTTPFunction(fn func(http.ResponseWriter, *http.Request)) (http.Handler
 			defer fmt.Println()
 			defer fmt.Fprintln(os.Stderr)
 		}
-		r, cancel := setContextTimeoutIfRequested(r)
+		r, cancel := setupRequestContext(r)
 		if cancel != nil {
 			defer cancel()
 		}
@@ -218,7 +218,7 @@ func wrapEventFunction(fn interface{}) (http.Handler, error) {
 			defer fmt.Println()
 			defer fmt.Fprintln(os.Stderr)
 		}
-		r, cancel := setContextTimeoutIfRequested(r)
+		r, cancel := setupRequestContext(r)
 		if cancel != nil {
 			defer cancel()
 		}
@@ -398,6 +398,12 @@ func writeHTTPErrorResponse(w http.ResponseWriter, statusCode int, status, msg s
 	fmt.Fprint(w, msg)
 }
 
+func setupRequestContext(r *http.Request) (*http.Request, func()) {
+	r, cancel := setContextTimeoutIfRequested(r)
+	r = addLoggingIDsToRequest(r)
+	return r, cancel
+}
+
 // setContextTimeoutIfRequested replaces the request's context with a cancellation if requested
 func setContextTimeoutIfRequested(r *http.Request) (*http.Request, func()) {
 	timeoutStr := os.Getenv("CLOUD_RUN_TIMEOUT_SECONDS")

funcframework/logging.go

@@ -0,0 +1,198 @@
+package funcframework
+
+import (
+	"bufio"
+	"context"
+	"encoding/json"
+	"io"
+	"net/http"
+	"os"
+	"regexp"
+	"sync"
+)
+
+var (
+	loggingIDsContextKey    contextKey = "loggingIDs"
+	validXCloudTraceContext            = regexp.MustCompile(
+		// Matches on "TRACE_ID"
+		`([a-f\d]+)?` +
+			// Matches on "/SPAN_ID"
+			`(?:/([a-f\d]+))?` +
+			// Matches on ";0=TRACE_TRUE"
+			`(?:;o=(\d))?`)
+)
+
+type loggingIDs struct {
+	trace       string
+	spanID      string
+	executionID string
+}
+
+type contextKey string
+
+func addLoggingIDsToRequest(r *http.Request) *http.Request {
+	executionID := r.Header.Get("Function-Execution-Id")
+	traceID, spanID, _ := deconstructXCloudTraceContext(r.Header.Get("X-Cloud-Trace-Context"))
+
+	if executionID == "" && traceID == "" && spanID == "" {
+		return r
+	}
+
+	r = r.WithContext(contextWithLoggingIDs(r.Context(), &loggingIDs{
+		trace:       traceID,
+		spanID:      spanID,
+		executionID: executionID,
+	}))
+
+	return r
+}
+
+func contextWithLoggingIDs(ctx context.Context, loggingIDs *loggingIDs) context.Context {
+	return context.WithValue(ctx, loggingIDsContextKey, loggingIDs)
+}
+
+func loggingIDsFromContext(ctx context.Context) *loggingIDs {
+	val := ctx.Value(loggingIDsContextKey)
+	if val == nil {
+		return nil
+	}
+	return val.(*loggingIDs)
+}
+
+func TraceIDFromContext(ctx context.Context) string {
+	ids := loggingIDsFromContext(ctx)
+	if ids == nil {
+		return ""
+	}
+	return ids.trace
+}
+
+func ExecutionIDFromContext(ctx context.Context) string {
+	ids := loggingIDsFromContext(ctx)
+	if ids == nil {
+		return ""
+	}
+	return ids.executionID
+}
+
+func SpanIDFromContext(ctx context.Context) string {
+	ids := loggingIDsFromContext(ctx)
+	if ids == nil {
+		return ""
+	}
+	return ids.spanID
+}
+
+func deconstructXCloudTraceContext(s string) (traceID, spanID string, traceSampled bool) {
+	// As per the format described at https://cloud.google.com/trace/docs/setup#force-trace
+	//    "X-Cloud-Trace-Context: TRACE_ID/SPAN_ID;o=TRACE_TRUE"
+	// for example:
+	//    "X-Cloud-Trace-Context: 105445aa7843bc8bf206b120001000/1;o=1"
+	matches := validXCloudTraceContext.FindStringSubmatch(s)
+	if matches != nil {
+		traceID, spanID, traceSampled = matches[1], matches[2], matches[3] == "1"
+	}
+	if spanID == "0" {
+		spanID = ""
+	}
+	return
+}
+
+// structuredLogEvent declares a subset of the fields supported by cloudlogging structured log events.
+// See https://cloud.google.com/logging/docs/structured-logging.
+type structuredLogEvent struct {
+	Message string            `json:"message"`
+	Trace   string            `json:"logging.googleapis.com/trace,omitempty"`
+	SpanID  string            `json:"logging.googleapis.com/spanId,omitempty"`
+	Labels  map[string]string `json:"logging.googleapis.com/labels,omitempty"`
+}
+
+// structuredLogWriter writes structured logs
+type structuredLogWriter struct {
+	mu         sync.Mutex
+	w          io.Writer
+	loggingIDs loggingIDs
+	buf        []byte
+}
+
+func (w *structuredLogWriter) writeStructuredLog(loggingIDs loggingIDs, message string) (int, error) {
+	event := structuredLogEvent{
+		Message: message,
+		Trace:   loggingIDs.trace,
+		SpanID:  loggingIDs.spanID,
+	}
+	if loggingIDs.executionID != "" {
+		event.Labels = map[string]string{
+			"execution_id": loggingIDs.executionID,
+		}
+	}
+
+	marshalled, err := json.Marshal(event)
+	if err != nil {
+		return 0, err
+	}
+	marshalled = append(marshalled, '\n')
+	return w.w.Write(marshalled)
+}
+
+func (w *structuredLogWriter) Write(output []byte) (int, error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+
+	w.buf = append(w.buf, output...)
+	buf := w.buf
+	wroteLines := 0
+	for {
+		advance, token, err := bufio.ScanLines(buf, false)
+		if token == nil || err != nil {
+			break
+		}
+		buf = buf[advance:]
+		if _, err := w.writeStructuredLog(w.loggingIDs, string(token)); err != nil {
+			return 0, err
+		}
+		wroteLines += 1
+	}
+
+	if wroteLines > 0 {
+		// Compact the buffer by copying remaining bytes to the start.
+		w.buf = append(w.buf[:0], buf...)
+	}
+
+	return len(output), nil
+}
+
+func (w *structuredLogWriter) Close() error {
+	if len(w.buf) == 0 {
+		return nil
+	}
+	_, err := w.writeStructuredLog(w.loggingIDs, string(w.buf))
+	return err
+}
+
+// LogWriter returns an io.Writer as a log sink for the request context.
+// One log event is generated for each new line terminated byte sequence
+// written to the io.Writer.
+//
+// This can be used with common logging frameworks, for example:
+//
+//	import (
+//	  "log"
+//	  "github.com/GoogleCloudPlatform/functions-framework-go/funcframework"
+//	)
+//	...
+//	func helloWorld(w http.ResponseWriter, r *http.Request) {
+//	  l := logger.New(funcframework.LogWriter(r.Context()))
+//	  l.Println("hello world!")
+//	}
+func LogWriter(ctx context.Context) io.WriteCloser {
+	loggingIDs := loggingIDsFromContext(ctx)
+	if loggingIDs == nil {
+		return os.Stderr
+	}
+
+	return &structuredLogWriter{
+		w:          os.Stderr,
+		loggingIDs: *loggingIDs,
+	}
+}

funcframework/logging_test.go

@@ -0,0 +1,135 @@
+package funcframework
+
+import (
+	"bytes"
+	"fmt"
+	"log"
+	"net/http/httptest"
+	"testing"
+)
+
+func TestLoggingIDExtraction(t *testing.T) {
+	tcs := []struct {
+		name            string
+		headers         map[string]string
+		wantTraceID     string
+		wantSpanID      string
+		wantExecutionID string
+	}{
+		{
+			name:    "no IDs",
+			headers: map[string]string{},
+		},
+		{
+			name: "execution ID only",
+			headers: map[string]string{
+				"Function-Execution-Id": "exec id",
+			},
+			wantExecutionID: "exec id",
+		},
+		{
+			name: "malformatted X-Cloud-Trace-Context",
+			headers: map[string]string{
+				"X-Cloud-Trace-Context": "$*#$(v434)",
+			},
+		},
+		{
+			name: "trace ID only",
+			headers: map[string]string{
+				"X-Cloud-Trace-Context": "0123456789abcdef",
+			},
+			wantTraceID: "0123456789abcdef",
+		},
+		{
+			name: "trace ID and span ID",
+			headers: map[string]string{
+				"X-Cloud-Trace-Context": "0123456789abcdef/aaaaaa",
+			},
+			wantTraceID: "0123456789abcdef",
+			wantSpanID:  "aaaaaa",
+		},
+		{
+			name: "all",
+			headers: map[string]string{
+				"X-Cloud-Trace-Context": "a/b",
+				"Function-Execution-Id": "c",
+			},
+			wantTraceID:     "a",
+			wantSpanID:      "b",
+			wantExecutionID: "c",
+		},
+	}
+
+	for _, tc := range tcs {
+		t.Run(tc.name, func(t *testing.T) {
+			r := httptest.NewRequest("POST", "/", bytes.NewReader(nil))
+			for k, v := range tc.headers {
+				r.Header.Set(k, v)
+			}
+			r = addLoggingIDsToRequest(r)
+			ctx := r.Context()
+
+			if tid := TraceIDFromContext(ctx); tid != tc.wantTraceID {
+				t.Errorf("expected trace id %q but got %q", tc.wantTraceID, tid)
+			}
+
+			if spid := SpanIDFromContext(ctx); spid != tc.wantSpanID {
+				t.Errorf("expected span id %q but got %q", tc.wantSpanID, spid)
+			}
+
+			if eid := ExecutionIDFromContext(ctx); eid != tc.wantExecutionID {
+				t.Errorf("expected execution id %q but got %q", tc.wantExecutionID, eid)
+			}
+		})
+	}
+}
+
+func TestStructuredLogWriter(t *testing.T) {
+	output := bytes.NewBuffer(nil)
+
+	w := &structuredLogWriter{
+		w: output,
+		loggingIDs: loggingIDs{
+			spanID:      "a",
+			trace:       "b",
+			executionID: "c",
+		},
+	}
+
+	fmt.Fprintf(w, "hello world!\n")
+	fmt.Fprintf(w, "this is another log line!\n")
+
+	wantOutput := `{"message":"hello world!","logging.googleapis.com/trace":"b","logging.googleapis.com/spanId":"a","logging.googleapis.com/labels":{"execution_id":"c"}}
+{"message":"this is another log line!","logging.googleapis.com/trace":"b","logging.googleapis.com/spanId":"a","logging.googleapis.com/labels":{"execution_id":"c"}}
+`
+	if output.String() != wantOutput {
+		t.Errorf("expected output %q got %q", wantOutput, output.String())
+	}
+}
+
+func TestLogPackageCompat(t *testing.T) {
+	output := bytes.NewBuffer(nil)
+	w := &structuredLogWriter{
+		w: output,
+		loggingIDs: loggingIDs{
+			spanID:      "a",
+			trace:       "b",
+			executionID: "c",
+		},
+	}
+
+	l := log.New(w, "", 0)
+	l.Print("go logger line")
+	l.Print("a second log line")
+	l.Print("a multiline\nstring in a single log\ncall")
+
+	wantOutput := `{"message":"go logger line","logging.googleapis.com/trace":"b","logging.googleapis.com/spanId":"a","logging.googleapis.com/labels":{"execution_id":"c"}}
+{"message":"a second log line","logging.googleapis.com/trace":"b","logging.googleapis.com/spanId":"a","logging.googleapis.com/labels":{"execution_id":"c"}}
+{"message":"a multiline","logging.googleapis.com/trace":"b","logging.googleapis.com/spanId":"a","logging.googleapis.com/labels":{"execution_id":"c"}}
+{"message":"string in a single log","logging.googleapis.com/trace":"b","logging.googleapis.com/spanId":"a","logging.googleapis.com/labels":{"execution_id":"c"}}
+{"message":"call","logging.googleapis.com/trace":"b","logging.googleapis.com/spanId":"a","logging.googleapis.com/labels":{"execution_id":"c"}}
+`
+	if output.String() != wantOutput {
+		t.Errorf("expected output %q got %q", wantOutput, output.String())
+	}
+}

go.mod

@@ -4,6 +4,7 @@ go 1.11
 
 require (
 	cloud.google.com/go/functions v1.16.2
+	cloud.google.com/go/logging v1.10.0 // indirect
 	github.com/cloudevents/sdk-go/v2 v2.14.0
 	github.com/google/go-cmp v0.6.0
 )