sdk/go/ctxlog/log.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: Apache-2.0
   4
   5 package ctxlog
   6
   7 import (
   8         "bytes"
   9         "context"
  10         "io"
  11         "os"
  12
  13         "github.com/sirupsen/logrus"
  14 )
  15
  16 var (
  17         loggerCtxKey = new(int)
  18         rootLogger   = logrus.New()
  19 )
  20
  21 const rfc3339NanoFixed = "2006-01-02T15:04:05.000000000Z07:00"
  22
  23 // Context returns a new child context such that FromContext(child)
  24 // returns the given logger.
  25 func Context(ctx context.Context, logger logrus.FieldLogger) context.Context {
  26         return context.WithValue(ctx, loggerCtxKey, logger)
  27 }
  28
  29 // FromContext returns the logger suitable for the given context -- the one
  30 // attached by contextWithLogger() if applicable, otherwise the
  31 // top-level logger with no fields/values.
  32 func FromContext(ctx context.Context) logrus.FieldLogger {
  33         if ctx != nil {
  34                 if logger, ok := ctx.Value(loggerCtxKey).(logrus.FieldLogger); ok {
  35                         return logger
  36                 }
  37         }
  38         return rootLogger.WithFields(nil)
  39 }
  40
  41 // New returns a new logger with the indicated format and
  42 // level.
  43 func New(out io.Writer, format, level string) *logrus.Logger {
  44         logger := logrus.New()
  45         logger.Out = out
  46         setFormat(logger, format)
  47         setLevel(logger, level)
  48         return logger
  49 }
  50
  51 func TestLogger(c interface{ Log(...interface{}) }) *logrus.Logger {
  52         logger := logrus.New()
  53         logger.Out = &logWriter{c.Log}
  54         setFormat(logger, "text")
  55         if d := os.Getenv("ARVADOS_DEBUG"); d != "0" && d != "" {
  56                 setLevel(logger, "debug")
  57         } else {
  58                 setLevel(logger, "info")
  59         }
  60         return logger
  61 }
  62
  63 // LogWriter returns an io.Writer that writes to the given log func,
  64 // which is typically (*check.C).Log().
  65 func LogWriter(log func(...interface{})) io.Writer {
  66         return &logWriter{log}
  67 }
  68
  69 // SetLevel sets the current logging level. See logrus for level
  70 // names.
  71 func SetLevel(level string) {
  72         setLevel(rootLogger, level)
  73 }
  74
  75 func setLevel(logger *logrus.Logger, level string) {
  76         if level == "" {
  77         } else if lvl, err := logrus.ParseLevel(level); err != nil {
  78                 logrus.WithField("Level", level).Fatal("unknown log level")
  79         } else {
  80                 logger.Level = lvl
  81         }
  82 }
  83
  84 // SetFormat sets the current logging format to "json" or "text".
  85 func SetFormat(format string) {
  86         setFormat(rootLogger, format)
  87 }
  88
  89 func setFormat(logger *logrus.Logger, format string) {
  90         switch format {
  91         case "text":
  92                 logger.Formatter = &logrus.TextFormatter{
  93                         FullTimestamp:   true,
  94                         TimestampFormat: rfc3339NanoFixed,
  95                 }
  96         case "json", "":
  97                 logger.Formatter = &logrus.JSONFormatter{
  98                         TimestampFormat: rfc3339NanoFixed,
  99                 }
 100         default:
 101                 logrus.WithField("Format", format).Fatal("unknown log format")
 102         }
 103 }
 104
 105 // logWriter is an io.Writer that writes by calling a "write log"
 106 // function, typically (*check.C)Log().
 107 type logWriter struct {
 108         logfunc func(...interface{})
 109 }
 110
 111 func (tl *logWriter) Write(buf []byte) (int, error) {
 112         tl.logfunc(string(bytes.TrimRight(buf, "\n")))
 113         return len(buf), nil
 114 }