logger.go

package logger

import (
	"context"
	"fmt"
	"net/http"
	"os"
	"strings"

	"github.com/getsentry/sentry-go"
	"github.com/vrischmann/envconfig"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
)

var global *zap.SugaredLogger

func init() {
	SetLogger(New(zap.NewAtomicLevelAt(zap.InfoLevel), nil))
}

// SetLogger sets global used logger. This function is not thread-safe.
func SetLogger(l *zap.SugaredLogger) {
	global = l
}

// New creates new *zap.SugaredLogger with standard EncoderConfig
func New(lvl zapcore.LevelEnabler, customFieldsKV []string, options ...zap.Option) *zap.SugaredLogger {
	sink := zapcore.AddSync(os.Stdout)
	options = append(options, zap.ErrorOutput(sink))
	cfg := zapcore.EncoderConfig{
		NameKey:        "logger",
		CallerKey:      "caller",
		StacktraceKey:  "stacktrace",
		MessageKey:     "message",
		LevelKey:       "severity",
		TimeKey:        "timestamp",
		LineEnding:     zapcore.DefaultLineEnding,
		EncodeTime:     zapcore.ISO8601TimeEncoder,
		EncodeDuration: zapcore.SecondsDurationEncoder,
		EncodeCaller:   zapcore.ShortCallerEncoder,
		EncodeLevel:    zapcore.LowercaseLevelEncoder,
	}
	encoder := zapcore.NewJSONEncoder(cfg)

	return zap.New(zapcore.NewCore(encoder, sink, lvl), options...).With(customFields(customFieldsKV)...).Sugar()
}

func customFields(customFieldsKV []string) []zapcore.Field {
	var fields []zapcore.Field

	for i := 0; i < len(customFieldsKV); i += 2 {
		fields = append(fields, zap.String(customFieldsKV[i], customFieldsKV[i+1]))
	}

	return fields
}

// Logger returns current global logger.
func Logger() *zap.SugaredLogger {
	return global
}

func InitLogger(serviceName string, customFields ...string) error {
	type Config struct {
		Level string `envconfig:"default=info"`
	}
	var cfg struct {
		Log *Config
	}
	if err := envconfig.InitWithPrefix(&cfg, serviceName); err != nil {
		return fmt.Errorf("cannot get log config; err: %w", err)
	}

	lvl := zap.NewAtomicLevel()
	err := lvl.UnmarshalText([]byte(cfg.Log.Level))
	if err != nil {
		return fmt.Errorf("failed to unmurshal log level: %s; err: %v", cfg.Log.Level, err)
	}

	SetLogger(New(lvl, customFields))

	return nil
}

type contextKey struct{}

var loggerContextKey = contextKey{}

// ToContext returns new context with specified sugared logger inside.
func ToContext(ctx context.Context, l *zap.SugaredLogger) context.Context {
	return context.WithValue(ctx, loggerContextKey, l)
}

// FromContext returns logger from context if set. Otherwise returns global `global` logger.
// In both cases returned logger is populated with `trace_id` & `span_id`.
func FromContext(ctx context.Context) *zap.SugaredLogger {
	var (
		logger *zap.SugaredLogger
		ok     bool
	)
	logger, ok = ctx.Value(loggerContextKey).(*zap.SugaredLogger)
	if !ok {
		logger = global
	}

	s, ok := tracer.SpanFromContext(ctx)
	if ok && s.Context().TraceID() > 0 {
		logger = withTraceID(logger, s.Context().TraceID(), s.Context().SpanID())
	}
	reqID := ctx.Value(requestIDKey{})
	if reqIDValue, ok := reqID.(string); ok {
		logger = withRequestID(logger, reqIDValue)
	}

	return logger
}

func withTraceID(l *zap.SugaredLogger, traceID, spanID uint64) *zap.SugaredLogger {
	return l.With(
		zap.Uint64("dd.trace_id", traceID),
		zap.Uint64("dd.span_id", spanID),
	)
}
func withRequestID(l *zap.SugaredLogger, requestID string) *zap.SugaredLogger {
	return l.With(
		zap.String("request_id", requestID),
	)
}

func Debug(ctx context.Context, args ...interface{}) {
	FromContext(ctx).Debug(args...)
}

func Debugf(ctx context.Context, format string, args ...interface{}) {
	FromContext(ctx).Debugf(format, args...)
}

func Info(ctx context.Context, args ...interface{}) {
	FromContext(ctx).Info(args...)
}

func Infof(ctx context.Context, format string, args ...interface{}) {
	FromContext(ctx).Infof(format, args...)
}

func Warn(ctx context.Context, args ...interface{}) {
	FromContext(ctx).Warn(args...)
}

func Warnf(ctx context.Context, format string, args ...interface{}) {
	FromContext(ctx).Warnf(format, args...)
}

func Error(ctx context.Context, args ...interface{}) {
	sendToSentry("", args...)
	markTracing(ctx, "", args...)
	FromContext(ctx).Error(args...)
}

func Errorf(ctx context.Context, format string, args ...interface{}) {
	sendToSentry(format, args...)
	markTracing(ctx, format, args...)
	FromContext(ctx).Errorf(format, args...)
}

func Fatal(ctx context.Context, args ...interface{}) {
	sendToSentry("", args...)
	markTracing(ctx, "", args...)
	FromContext(ctx).Fatal(args...)
}

func Fatalf(ctx context.Context, format string, args ...interface{}) {
	sendToSentry(format, args...)
	markTracing(ctx, format, args...)
	FromContext(ctx).Fatalf(format, args...)
}

func prepareErrFormat(format string, args ...interface{}) string {
	if format != "" || len(args) == 0 {
		return format
	}

	return strings.Repeat("%+v ", len(args))
}
func sendToSentry(format string, args ...interface{}) {
	format = prepareErrFormat(format, args)
	if format != "" {
		sentry.CaptureException(fmt.Errorf(format, args...))
	}
}

func markTracing(ctx context.Context, format string, args ...interface{}) {
	span, ok := tracer.SpanFromContext(ctx)
	if !ok {
		return
	}

	format = prepareErrFormat(format, args)
	if format != "" {
		span.SetTag("error", fmt.Errorf(format, args...))
	}
}

type requestIDKey struct{}

func ReqWithLoggerContext(r *http.Request) *http.Request {
	reqID := r.Header.Get("X-Request-Id")
	if reqID != "" {
		return r.WithContext(context.WithValue(r.Context(), requestIDKey{}, reqID))
	}

	return r
}

func AddHeadersFromContext(headers http.Header, ctx context.Context) {
	reqIDCtx := ctx.Value(requestIDKey{})
	if reqIDCtx == nil {
		return
	}

	reqID, ok := reqIDCtx.(string)
	if !ok {
		return
	}
	headers.Set("X-Request-Id", reqID)
}