writer.go 3.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102
  1. package logrus
  2. import (
  3. "bufio"
  4. "io"
  5. "runtime"
  6. "strings"
  7. )
  8. // Writer at INFO level. See WriterLevel for details.
  9. func (logger *Logger) Writer() *io.PipeWriter {
  10. return logger.WriterLevel(InfoLevel)
  11. }
  12. // WriterLevel returns an io.Writer that can be used to write arbitrary text to
  13. // the logger at the given log level. Each line written to the writer will be
  14. // printed in the usual way using formatters and hooks. The writer is part of an
  15. // io.Pipe and it is the callers responsibility to close the writer when done.
  16. // This can be used to override the standard library logger easily.
  17. func (logger *Logger) WriterLevel(level Level) *io.PipeWriter {
  18. return NewEntry(logger).WriterLevel(level)
  19. }
  20. // Writer returns an io.Writer that writes to the logger at the info log level
  21. func (entry *Entry) Writer() *io.PipeWriter {
  22. return entry.WriterLevel(InfoLevel)
  23. }
  24. // WriterLevel returns an io.Writer that writes to the logger at the given log level
  25. func (entry *Entry) WriterLevel(level Level) *io.PipeWriter {
  26. reader, writer := io.Pipe()
  27. var printFunc func(args ...interface{})
  28. // Determine which log function to use based on the specified log level
  29. switch level {
  30. case TraceLevel:
  31. printFunc = entry.Trace
  32. case DebugLevel:
  33. printFunc = entry.Debug
  34. case InfoLevel:
  35. printFunc = entry.Info
  36. case WarnLevel:
  37. printFunc = entry.Warn
  38. case ErrorLevel:
  39. printFunc = entry.Error
  40. case FatalLevel:
  41. printFunc = entry.Fatal
  42. case PanicLevel:
  43. printFunc = entry.Panic
  44. default:
  45. printFunc = entry.Print
  46. }
  47. // Start a new goroutine to scan the input and write it to the logger using the specified print function.
  48. // It splits the input into chunks of up to 64KB to avoid buffer overflows.
  49. go entry.writerScanner(reader, printFunc)
  50. // Set a finalizer function to close the writer when it is garbage collected
  51. runtime.SetFinalizer(writer, writerFinalizer)
  52. return writer
  53. }
  54. // writerScanner scans the input from the reader and writes it to the logger
  55. func (entry *Entry) writerScanner(reader *io.PipeReader, printFunc func(args ...interface{})) {
  56. scanner := bufio.NewScanner(reader)
  57. // Set the buffer size to the maximum token size to avoid buffer overflows
  58. scanner.Buffer(make([]byte, bufio.MaxScanTokenSize), bufio.MaxScanTokenSize)
  59. // Define a split function to split the input into chunks of up to 64KB
  60. chunkSize := bufio.MaxScanTokenSize // 64KB
  61. splitFunc := func(data []byte, atEOF bool) (int, []byte, error) {
  62. if len(data) >= chunkSize {
  63. return chunkSize, data[:chunkSize], nil
  64. }
  65. return bufio.ScanLines(data, atEOF)
  66. }
  67. // Use the custom split function to split the input
  68. scanner.Split(splitFunc)
  69. // Scan the input and write it to the logger using the specified print function
  70. for scanner.Scan() {
  71. printFunc(strings.TrimRight(scanner.Text(), "\r\n"))
  72. }
  73. // If there was an error while scanning the input, log an error
  74. if err := scanner.Err(); err != nil {
  75. entry.Errorf("Error while reading from Writer: %s", err)
  76. }
  77. // Close the reader when we are done
  78. reader.Close()
  79. }
  80. // WriterFinalizer is a finalizer function that closes then given writer when it is garbage collected
  81. func writerFinalizer(writer *io.PipeWriter) {
  82. writer.Close()
  83. }