logging/logging.go

package logging

import (
	"context"
	"os"
	"path"

	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
)

type loggerKeyType int

const loggerKey loggerKeyType = iota

var defaultLogger *zap.Logger

func init() {
	l, _ := zap.NewDevelopment()
	defaultLogger = l.With(
		zap.Int("pid", os.Getpid()),
		zap.String("exe", path.Base(os.Args[0])),
	)
}

func NewContext(ctx context.Context, logger *zap.Logger, fields ...zapcore.Field) context.Context {
	if nil != logger {
		return context.WithValue(ctx, loggerKey, logger.With(fields...))
	}
	return context.WithValue(ctx, loggerKey, WithContext(ctx).With(fields...))
}

func WithContext(ctx context.Context) *zap.Logger {
	if ctx == nil {
		return defaultLogger
	}
	if ctxLogger, ok := ctx.Value(loggerKey).(*zap.Logger); ok {
		return ctxLogger
	}
	return defaultLogger
}