github.com/yandex/pandora@v0.5.32/lib/zaputil/stack_extract_core.go (about)

     1  package zaputil
     2  
     3  import (
     4  	"fmt"
     5  
     6  	"github.com/pkg/errors"
     7  	"go.uber.org/zap"
     8  	"go.uber.org/zap/buffer"
     9  	"go.uber.org/zap/zapcore"
    10  )
    11  
    12  // NewStackExtractCore returns core that extracts stacktraces from
    13  // error fields with github.com/pkg/errors error types, and append them to
    14  // zapcore.Entry.Stack, on Write.
    15  // That makes stacktraces from errors readable in case of console encoder.
    16  // WARN(skipor): don't call Check of underlying cores, just use LevelEnabler.
    17  // That breaks sampling and other complex logic of choosing log or not entry.
    18  func NewStackExtractCore(c zapcore.Core) zapcore.Core {
    19  	return &errStackExtractCore{c, getBuffer()}
    20  }
    21  
    22  type errStackExtractCore struct {
    23  	zapcore.Core
    24  	stacksBuff zapBuffer
    25  }
    26  
    27  type stackedErr interface {
    28  	error
    29  	StackTrace() errors.StackTrace
    30  }
    31  
    32  type causer interface {
    33  	Cause() error
    34  }
    35  
    36  func (c *errStackExtractCore) With(fields []zapcore.Field) zapcore.Core {
    37  	buff := c.cloneBuffer()
    38  	fields = extractFieldsStacksToBuff(buff, fields)
    39  	return &errStackExtractCore{
    40  		c.Core.With(fields),
    41  		buff,
    42  	}
    43  }
    44  
    45  func (c *errStackExtractCore) Write(ent zapcore.Entry, fields []zapcore.Field) error {
    46  	if c.stacksBuff.Len() == 0 && !hasStacksToExtract(fields) {
    47  		return c.Core.Write(ent, fields)
    48  	}
    49  	buff := c.cloneBuffer()
    50  	defer buff.Free()
    51  	fields = extractFieldsStacksToBuff(buff, fields)
    52  
    53  	if ent.Stack == "" {
    54  		ent.Stack = buff.String()
    55  	} else {
    56  		// Should be rare case, so allocation is OK.
    57  		ent.Stack = ent.Stack + "\n" + buff.String()
    58  	}
    59  	return c.Core.Write(ent, fields)
    60  }
    61  
    62  func (c *errStackExtractCore) Check(ent zapcore.Entry, ce *zapcore.CheckedEntry) *zapcore.CheckedEntry {
    63  	// HACK(skipor): not calling Check of nested. It's ok while we use simple io/tee cores.
    64  	// But that breaks sampling logic of underlying cores, for example.
    65  	if c.Enabled(ent.Level) {
    66  		return ce.AddCore(ent, c)
    67  	}
    68  	return ce
    69  }
    70  
    71  func (c *errStackExtractCore) cloneBuffer() zapBuffer {
    72  	clone := getBuffer()
    73  	_, _ = clone.Write(c.stacksBuff.Bytes())
    74  	return clone
    75  }
    76  
    77  func hasStacksToExtract(fields []zapcore.Field) bool {
    78  	for _, field := range fields {
    79  		if field.Type != zapcore.ErrorType {
    80  			continue
    81  		}
    82  		_, ok := field.Interface.(stackedErr)
    83  		if ok {
    84  			return true
    85  		}
    86  	}
    87  	return false
    88  }
    89  
    90  func extractFieldsStacksToBuff(buff zapBuffer, fields []zapcore.Field) []zapcore.Field {
    91  	var stacksFound bool
    92  	for i, field := range fields {
    93  		if field.Type != zapcore.ErrorType {
    94  			continue
    95  		}
    96  		stacked, ok := field.Interface.(stackedErr)
    97  		if !ok {
    98  			continue
    99  		}
   100  		if !stacksFound {
   101  			stacksFound = true
   102  			oldFields := fields
   103  			fields = make([]zapcore.Field, len(fields))
   104  			copy(fields, oldFields)
   105  		}
   106  		if cause, ok := stacked.(causer); ok {
   107  			field.Interface = cause.Cause()
   108  		} else {
   109  			field = zap.String(field.Key, stacked.Error())
   110  		}
   111  		fields[i] = field
   112  		appendStack(buff, field.Key, stacked.StackTrace())
   113  	}
   114  	return fields // Cloned in case modifications.
   115  }
   116  
   117  func appendStack(buff zapBuffer, key string, stack errors.StackTrace) {
   118  	if buff.Len() != 0 {
   119  		buff.AppendByte('\n')
   120  	}
   121  	buff.AppendString(key)
   122  	buff.AppendString(" stacktrace:")
   123  	stack.Format(zapBufferFmtState{buff}, 'v')
   124  }
   125  
   126  type zapBuffer struct{ *buffer.Buffer }
   127  
   128  var _ ioStringWriter = zapBuffer{}
   129  
   130  type ioStringWriter interface {
   131  	WriteString(s string) (n int, err error)
   132  }
   133  
   134  func (b zapBuffer) WriteString(s string) (n int, err error) {
   135  	b.AppendString(s)
   136  	return len(s), nil
   137  }
   138  
   139  var bufferPool = buffer.NewPool()
   140  
   141  func getBuffer() zapBuffer {
   142  	return zapBuffer{bufferPool.Get()}
   143  }
   144  
   145  type zapBufferFmtState struct{ zapBuffer }
   146  
   147  var _ fmt.State = zapBufferFmtState{}
   148  
   149  func (zapBufferFmtState) Flag(c int) bool {
   150  	switch c {
   151  	case '+':
   152  		return true
   153  	default:
   154  		return false
   155  	}
   156  }
   157  
   158  func (zapBufferFmtState) Width() (wid int, ok bool)      { panic("should not be called") }
   159  func (zapBufferFmtState) Precision() (prec int, ok bool) { panic("should not be called") }