« Back to Index

Datadog Structured Logging Abstraction with Logrus

View original Gist on GitHub

Tags: #go #logs

logging.go

package logging

import (
	"fmt"
	"math/rand"
	"runtime/debug"

	"github.com/sirupsen/logrus"
	"github.com/sirupsen/logrus/hooks/test"
)

const (
	// TimestampFormat is the default timestamp format for log records:
	// https://golang.org/pkg/time/#Time.Format
	TimestampFormat = "2006-01-02 15:04:05.123"

	// Map these constants to logrus for the sake of SetLevel() function.
	PanicLevel = logrus.PanicLevel
	FatalLevel = logrus.FatalLevel
	ErrorLevel = logrus.ErrorLevel
	WarnLevel  = logrus.WarnLevel
	InfoLevel  = logrus.InfoLevel
	DebugLevel = logrus.DebugLevel
	TraceLevel = logrus.TraceLevel
)

// Fields is an alias for a set of key-value pairs to be logged
type Fields map[string]interface{}

// extend updates a set of Fields using only the values in "fields" which are also in "fieldNames"
func (s Fields) extend(fieldNames []string, fields Fields) {
	for _, fieldName := range fieldNames {
		if val, ok := fields[fieldName]; ok {
			s[fieldName] = val
		}
	}
}

// Logger is a logrus entry wrapped with functions to support structured logging
type Logger struct {
	raw    *logrus.Logger
	logger *logrus.Entry
	Rate   int
}

// NewLogger creates a new Logger instance
func NewLogger(service string, optFuncs ...func(*logrus.Logger)) *Logger {
	logger := logrus.New()
	logger.SetFormatter(&logrus.JSONFormatter{})

	for _, f := range optFuncs {
		f(logger)
	}

	entry := logger.WithField("service", service)

	// initialize logging
	return &Logger{
		raw:    logger,
		logger: entry,
		Rate:   10,
	}
}

// TestLogger returns a logger that can be used for testing
func TestLogger(service string, optFuncs ...func(*logrus.Logger)) (*Logger, *test.Hook) {
	logger, hook := test.NewNullLogger()
	logger.SetFormatter(&logrus.JSONFormatter{})

	entry := logger.WithField("service", service)

	return &Logger{
		logger: entry,
		Rate:   10,
	}, hook
}

// getField returns a single field in the logs
func (l *Logger) getField(key string) (interface{}, bool) {
	if field, ok := l.logger.Data[key]; ok {
		return field, ok
	}
	return nil, false
}

// updateFields sets fields in a named set of fields, reusing existing values if possible
func (l *Logger) updateFields(key string, fieldNames []string, fields Fields) *Logger {
	var fieldSet Fields

	// get the existing set of fields or create a new one
	fieldSet, ok := l.Fields(key)
	if !ok {
		fieldSet = Fields{}
	}

	// update the set of fields
	fieldSet.extend(fieldNames, fields)
	return &Logger{
		logger: l.logger.WithField(key, fieldSet),
		Rate:   l.Rate,
	}
}

// appendFields appends a Fields object to an array of Fields objects
func (l *Logger) appendFields(key string, fieldNames []string, fields Fields) *Logger {
	var fieldsArray []Fields
	filtered := Fields{}

	// get the existing array or create a new one
	if field, ok := l.getField(key); ok {
		fieldsArray, ok = field.([]Fields)
		if !ok {
			fieldsArray = []Fields{}
		}
	}

	filtered.extend(fieldNames, fields)

	// append the stanza to the array
	fieldsArray = append(fieldsArray, filtered)
	return &Logger{
		logger: l.logger.WithField(key, fieldsArray),
		Rate:   l.Rate,
	}
}

// Fields returns a set of fields in the logger if it exists
func (l *Logger) Fields(key string) (Fields, bool) {
	var fields Fields
	var data interface{}
	var ok bool

	// make sure the set exists
	if data, ok = l.getField(key); !ok {
		return nil, false
	}

	// make sure the set isn't malformed
	fields, ok = data.(Fields)
	if !ok {
		return nil, false
	}
	return fields, true
}

// Data returns the raw structured logging data from the logrus logger
func (l *Logger) Data() map[string]interface{} {
	return l.logger.Data
}

// sample returns True randomly at a percentage of the time
func sample(percent int) bool {
	return rand.Intn(100) <= percent
}

// Wrappers around logging actions

// InfoSampled wraps the logrus Info function and samples the logs
func (l *Logger) InfoSampled(args ...interface{}) {
	// we check the Rate==100 first to avoid the more expensive call to the sample rand function
	// the global rand is probably fine, but just to note that using it could slow things
	// down a bit since every call to it locks a mutex
	if l.Rate == 100 || l.logger.Level >= logrus.DebugLevel || sample(l.Rate) { // shhhh
		l.logger.Info(args...)
	}
}

// Info wraps the logrus Info function
func (l *Logger) Info(args ...interface{}) {
	l.logger.Info(args...)
}

// Warn logs a message at warning level
func (l *Logger) Warn(args ...interface{}) *Logger {
	l.logger.Warn(args...)
	return l
}

// Error logs a message at error level
func (l *Logger) Error(args ...interface{}) *Logger {
	l.logger.Error(args...)
	return l
}

// Fatal wraps the logrus Fatal function
func (l *Logger) Fatal(args ...interface{}) *Logger {
	l.logger.Fatal(args...)
	return l
}

// Debug wraps the logrus Debug function
func (l *Logger) Debug(args ...interface{}) *Logger {
	l.logger.Debug(args...)
	return l
}

// Panic wraps the logrus Panic function
func (l *Logger) Panic(args ...interface{}) *Logger {
	l.logger.Panic(args...)
	return l
}

// Printf wraps the logrus Printf function
func (l *Logger) Printf(format string, args ...interface{}) {
	l.logger.Printf(format, args...)
}

// Wrappers for common fields

// WithError adds error information to the logging record
func (l *Logger) WithError(err error) *Logger {
	return l.WithErrorFields(Fields{
		"message": err.Error(),
		"kind":    fmt.Sprintf("%T", err),
	})
}

// WithStack adds a stack trace to the logging record
func (l *Logger) WithStack() *Logger {
	return l.WithErrorFields(Fields{
		"stack": debug.Stack(),
	})
}

// WithFields sets key-value pairs in the "meta" set. Use `l.logger.WithField`
// in order to set a field directly in the logrus.Entry.
func (l *Logger) WithFields(fields Fields) *Logger {
	// get the existing set of fields or create a new one
	metaFields, ok := l.Fields("meta")
	if !ok {
		metaFields = Fields{}
	}

	// we skip the filtering check in withFields because anything goes
	// in the "meta" fields
	for key, val := range fields {
		metaFields[key] = val
	}

	return &Logger{
		logger: l.logger.WithField("meta", metaFields),
		Rate:   l.Rate,
	}
}

// WithField is a helper to add a single field to the meta set
func (l *Logger) WithField(key string, value interface{}) *Logger {
	return l.WithFields(Fields{key: value})
}

// WithBasicFields is a helper to set top level fields like timestamps
// This matches the `RootFieldSet` in `bf_storage.logging` py lib,
// however it has not been renamed yet since some go instrumentation relies on it already.
func (l *Logger) WithBasicFields(fields Fields) *Logger {
	var raw map[string]interface{}
	basicFields := Fields{}
	fieldNames := []string{
		"cluster",
		"level",
		"message",
		"region",
		"service",
		"source",
	}

	basicFields.extend(fieldNames, fields)
	raw = basicFields
	return &Logger{
		logger: l.logger.WithFields(raw),
		Rate:   l.Rate,
	}
}

// WithDatabaseFields replicates the `bf_storage.logging` DatabaseFieldSet.
// This is in line with the BuzzFeed established standard.
func (l *Logger) WithDatabaseFields(fields Fields) *Logger {
	fieldNames := []string{
		"instance",
		"statement",
		"operation",
		"user",
	}
	return l.updateFields("database", fieldNames, fields)
}

// WithErrorFields sets fields in the "error" set.
// We use Datadog's recommendation of "stack", "message", and "field".
func (l *Logger) WithErrorFields(fields Fields) *Logger {
	fieldNames := []string{
		"stack",
		"message",
		"kind",
	}
	return l.updateFields("error", fieldNames, fields)
}

// WithHTTPFields sets fields in the "http" set
// This fieldset uses the DataDog standard naming conventions
// https://docs.datadoghq.com/logs/processing/attributes_naming_convention/#http-requests
func (l *Logger) WithHTTPFields(fields Fields) *Logger {
	fieldNames := []string{
		"url",
		"method",
		"status_code",
		"referer",
		"request_id",
		"useragent",
		"request_time",
		"duration",
	}
	return l.updateFields("http", fieldNames, fields)
}

// WithLoggerFields supplies the standard logger fieldset.
// See DataDog docs here for naming conventions:
// https://docs.datadoghq.com/logs/processing/attributes_naming_convention/#source-code
func (l *Logger) WithLoggerFields(fields Fields) *Logger {
	fieldNames := []string{
		"name",
		"thread_name",
		"path_name",
		"method_name",
		"module",
		"line",
		"level_name",
		"version",
	}
	return l.updateFields("logger", fieldNames, fields)
}

// WithRigFields sets fields in the "rig" set
// This is used internally to add rig env data to each entry.
func (l *Logger) WithRigFields(fields Fields) *Logger {
	fieldNames := []string{
		"service",
		"cluster",
		"version",
		"pid",
		"region",
	}
	return l.updateFields("rig", fieldNames, fields)
}

// WithStorageFields replicates the `bf_storage.logging` StorageFieldSet.
// This is in line with the BuzzFeed established standard.
func (l *Logger) WithStorageFields(fields Fields) *Logger {
	fieldNames := []string{
		"bucket_name",
		"redis_hostname",
		"redis_isolated_hostname",
		"redis_shared_hostname",
		"max_thread_count",
		"mysql_database_name",
		"mysql_port",
		"mysql_port",
		"mysql_hostname",
		"mysql_isolated_hostname",
		"mysql_shared_hostname",
		"mysql_isolated_hostname",
		"mysql_username",
		"mysql_isolated_database_name",
		"mysql_shared_database_name",
		"environment_type",
		"version",
	}
	return l.updateFields("storage", fieldNames, fields)
}

// WithUpstreamFields replicates `UpstreamFieldSet` in `bf_storage.logging`.
// This brings the upstream fieldset inline with the BuzzFeed convention.
func (l *Logger) WithUpstreamFields(fields Fields) *Logger {
	fieldNames := []string{
		"service",
		"request_url",
		"status_code",
		"response_time",
	}
	return l.updateFields("upstream", fieldNames, fields)
}

// SetLevel sets the default log level on the underlying logrus instance.
//
func (l *Logger) SetLevel(level logrus.Level) {
	l.raw.Level = level
}

logging_test.go

package logging

import (
	"bytes"
	"testing"

	"github.com/sirupsen/logrus"
)

type modifier func(*logrus.Logger)

func modifyLogrusOutput(b *bytes.Buffer) modifier {
	return func(l *logrus.Logger) {
		l.SetOutput(b)
	}
}

func modifyLogrusLevel(l *logrus.Logger) {
	l.Level = logrus.DebugLevel
}

// TestDebugLogLevelSetDirectly validates that a consumer can change settings
// directly on the underlying logrus instance, such as the debug log level.
//
// NOTE: this requires the consumer to import the logrus package as a
// dependency inside their service.
//
func TestDebugLogLevelSetDirectly(t *testing.T) {
	var b bytes.Buffer

	log := NewLogger("golang_logging", modifyLogrusOutput(&b), modifyLogrusLevel)
	log.Debug("debug!")
	log.Info("info!")

	if !bytes.Contains(b.Bytes(), []byte("debug!")) {
		t.Error("debug level wasn't set")
	}

	if !bytes.Contains(b.Bytes(), []byte("info!")) {
		t.Error("info level wasn't set")
	}
}

// TestDebugLogLevelSetIndirectly validates that a consumer of the logging
// library doesn't have to add logrus as a dependency if they only need to set
// the log level (which is a standard feature consumers are likely to use).
//
// NOTE: for the sake of the test environment I still need to import logrus and
// modify it 'directly' in order for it to send log output to a buffer instead
// of to stdout, but I should be able to still set the log level 'indirectly'.
//
func TestDebugLogLevelSetIndirectly(t *testing.T) {
	var b bytes.Buffer

	log := NewLogger("golang_logging", modifyLogrusOutput(&b))
	log.SetLevel(DebugLevel)

	log.Debug("debug!")
	log.Info("info!")

	if !bytes.Contains(b.Bytes(), []byte("debug!")) {
		t.Error("debug level wasn't set")
	}

	if !bytes.Contains(b.Bytes(), []byte("info!")) {
		t.Error("info level wasn't set")
	}
}