diff --git a/log/context.go b/log/context.go new file mode 100644 index 0000000..4539e47 --- /dev/null +++ b/log/context.go @@ -0,0 +1,81 @@ +package log + +import ( + "path" + + "github.com/Sirupsen/logrus" + "golang.org/x/net/context" +) + +var ( + // G is an alias for GetLogger. + // + // We may want to define this locally to a package to get package tagged log + // messages. + G = GetLogger + + // L is an alias for the the standard logger. + L = logrus.NewEntry(logrus.StandardLogger()) +) + +type ( + loggerKey struct{} + moduleKey struct{} +) + +// WithLogger returns a new context with the provided logger. Use in +// combination with logger.WithField(s) for great effect. +func WithLogger(ctx context.Context, logger *logrus.Entry) context.Context { + return context.WithValue(ctx, loggerKey{}, logger) +} + +// GetLogger retrieves the current logger from the context. If no logger is +// available, the default logger is returned. +func GetLogger(ctx context.Context) *logrus.Entry { + logger := ctx.Value(loggerKey{}) + + if logger == nil { + return L + } + + return logger.(*logrus.Entry) +} + +// WithModule adds the module to the context, appending it with a slash if a +// module already exists. A module is just an roughly correlated defined by the +// call tree for a given context. +// +// As an example, we might have a "node" module already part of a context. If +// this function is called with "tls", the new value of module will be +// "node/tls". +// +// Modules represent the call path. If the new module and last module are the +// same, a new module entry will not be created. If the new module and old +// older module are the same but separated by other modules, the cycle will be +// represented by the module path. +func WithModule(ctx context.Context, module string) context.Context { + parent := GetModulePath(ctx) + + if parent != "" { + // don't re-append module when module is the same. + if path.Base(parent) == module { + return ctx + } + + module = path.Join(parent, module) + } + + ctx = WithLogger(ctx, GetLogger(ctx).WithField("module", module)) + return context.WithValue(ctx, moduleKey{}, module) +} + +// GetModulePath returns the module path for the provided context. If no module +// is set, an empty string is returned. +func GetModulePath(ctx context.Context) string { + module := ctx.Value(moduleKey{}) + if module == nil { + return "" + } + + return module.(string) +} diff --git a/log/context_test.go b/log/context_test.go new file mode 100644 index 0000000..ddff398 --- /dev/null +++ b/log/context_test.go @@ -0,0 +1,41 @@ +package log + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "golang.org/x/net/context" +) + +func TestLoggerContext(t *testing.T) { + ctx := context.Background() + assert.Equal(t, GetLogger(ctx), L) // should be same as L variable + assert.Equal(t, G(ctx), GetLogger(ctx)) // these should be the same. + + ctx = WithLogger(ctx, G(ctx).WithField("test", "one")) + assert.Equal(t, GetLogger(ctx).Data["test"], "one") + assert.Equal(t, G(ctx), GetLogger(ctx)) // these should be the same. +} + +func TestModuleContext(t *testing.T) { + ctx := context.Background() + assert.Equal(t, GetModulePath(ctx), "") + + ctx = WithModule(ctx, "a") // basic behavior + assert.Equal(t, GetModulePath(ctx), "a") + logger := GetLogger(ctx) + assert.Equal(t, logger.Data["module"], "a") + + parent, ctx := ctx, WithModule(ctx, "a") + assert.Equal(t, ctx, parent) // should be a no-op + assert.Equal(t, GetModulePath(ctx), "a") + assert.Equal(t, GetLogger(ctx).Data["module"], "a") + + ctx = WithModule(ctx, "b") // new module + assert.Equal(t, GetModulePath(ctx), "a/b") + assert.Equal(t, GetLogger(ctx).Data["module"], "a/b") + + ctx = WithModule(ctx, "c") // new module + assert.Equal(t, GetModulePath(ctx), "a/b/c") + assert.Equal(t, GetLogger(ctx).Data["module"], "a/b/c") +} diff --git a/log/grpc.go b/log/grpc.go new file mode 100644 index 0000000..4978d49 --- /dev/null +++ b/log/grpc.go @@ -0,0 +1,13 @@ +package log + +import ( + "golang.org/x/net/context" + "google.golang.org/grpc/grpclog" +) + +func init() { + ctx := WithModule(context.Background(), "grpc") + + // completely replace the grpc logger with the logrus logger. + grpclog.SetLogger(G(ctx)) +}