trpc.group/trpc-go/trpc-go@v1.0.2/log/rollwriter/async_roll_writer.go (about)

     1  //
     2  //
     3  // Tencent is pleased to support the open source community by making tRPC available.
     4  //
     5  // Copyright (C) 2023 THL A29 Limited, a Tencent company.
     6  // All rights reserved.
     7  //
     8  // If you have downloaded a copy of the tRPC source code from Tencent,
     9  // please note that tRPC source code is licensed under the  Apache 2.0 License,
    10  // A copy of the Apache 2.0 License is included in this file.
    11  //
    12  //
    13  
    14  package rollwriter
    15  
    16  import (
    17  	"bytes"
    18  	"errors"
    19  	"fmt"
    20  	"io"
    21  	"time"
    22  
    23  	"github.com/hashicorp/go-multierror"
    24  
    25  	"trpc.group/trpc-go/trpc-go/internal/report"
    26  )
    27  
    28  const (
    29  	defaultLogQueueSize    = 10000
    30  	defaultWriteLogSize    = 4 * 1024 // 4KB
    31  	defaultLogIntervalInMs = 100
    32  	defaultDropLog         = false
    33  )
    34  
    35  // AsyncRollWriter is the asynchronous rolling log writer which implements zapcore.WriteSyncer.
    36  type AsyncRollWriter struct {
    37  	logger io.WriteCloser
    38  	opts   *AsyncOptions
    39  
    40  	logQueue chan []byte
    41  	sync     chan struct{}
    42  	syncErr  chan error
    43  	close    chan struct{}
    44  	closeErr chan error
    45  }
    46  
    47  // NewAsyncRollWriter creates a new AsyncRollWriter.
    48  func NewAsyncRollWriter(logger io.WriteCloser, opt ...AsyncOption) *AsyncRollWriter {
    49  	opts := &AsyncOptions{
    50  		LogQueueSize:     defaultLogQueueSize,
    51  		WriteLogSize:     defaultWriteLogSize,
    52  		WriteLogInterval: defaultLogIntervalInMs,
    53  		DropLog:          defaultDropLog,
    54  	}
    55  
    56  	for _, o := range opt {
    57  		o(opts)
    58  	}
    59  
    60  	w := &AsyncRollWriter{
    61  		logger:   logger,
    62  		opts:     opts,
    63  		logQueue: make(chan []byte, opts.LogQueueSize),
    64  		sync:     make(chan struct{}),
    65  		syncErr:  make(chan error),
    66  		close:    make(chan struct{}),
    67  		closeErr: make(chan error),
    68  	}
    69  
    70  	// Start a new goroutine to write batch logs.
    71  	go w.batchWriteLog()
    72  	return w
    73  }
    74  
    75  // Write writes logs. It implements io.Writer.
    76  func (w *AsyncRollWriter) Write(data []byte) (int, error) {
    77  	log := make([]byte, len(data))
    78  	copy(log, data)
    79  	if w.opts.DropLog {
    80  		select {
    81  		case w.logQueue <- log:
    82  		default:
    83  			report.LogQueueDropNum.Incr()
    84  			return 0, errors.New("async roll writer: log queue is full")
    85  		}
    86  		return len(data), nil
    87  	}
    88  	w.logQueue <- log
    89  	return len(data), nil
    90  }
    91  
    92  // Sync syncs logs. It implements zapcore.WriteSyncer.
    93  func (w *AsyncRollWriter) Sync() error {
    94  	w.sync <- struct{}{}
    95  	return <-w.syncErr
    96  }
    97  
    98  // Close closes current log file. It implements io.Closer.
    99  func (w *AsyncRollWriter) Close() error {
   100  	err := w.Sync()
   101  	close(w.close)
   102  	return multierror.Append(err, <-w.closeErr).ErrorOrNil()
   103  }
   104  
   105  // batchWriteLog asynchronously writes logs in batches.
   106  func (w *AsyncRollWriter) batchWriteLog() {
   107  	buffer := bytes.NewBuffer(make([]byte, 0, w.opts.WriteLogSize*2))
   108  	ticker := time.NewTicker(time.Millisecond * time.Duration(w.opts.WriteLogInterval))
   109  	defer ticker.Stop()
   110  
   111  	for {
   112  		select {
   113  		case <-ticker.C:
   114  			if buffer.Len() > 0 {
   115  				_, err := w.logger.Write(buffer.Bytes())
   116  				handleErr(err, "w.logger.Write on tick")
   117  				buffer.Reset()
   118  			}
   119  		case data := <-w.logQueue:
   120  			buffer.Write(data)
   121  			if buffer.Len() >= w.opts.WriteLogSize {
   122  				_, err := w.logger.Write(buffer.Bytes())
   123  				handleErr(err, "w.logger.Write on log queue")
   124  				buffer.Reset()
   125  			}
   126  		case <-w.sync:
   127  			var err error
   128  			if buffer.Len() > 0 {
   129  				_, e := w.logger.Write(buffer.Bytes())
   130  				err = multierror.Append(err, e).ErrorOrNil()
   131  				buffer.Reset()
   132  			}
   133  			size := len(w.logQueue)
   134  			for i := 0; i < size; i++ {
   135  				v := <-w.logQueue
   136  				_, e := w.logger.Write(v)
   137  				err = multierror.Append(err, e).ErrorOrNil()
   138  			}
   139  			w.syncErr <- err
   140  		case <-w.close:
   141  			w.closeErr <- w.logger.Close()
   142  			return
   143  		}
   144  	}
   145  }
   146  
   147  func handleErr(err error, msg string) {
   148  	if err == nil {
   149  		return
   150  	}
   151  	// Log writer has errors, so output to stdout directly.
   152  	fmt.Printf("async roll writer err: %+v, msg: %s", err, msg)
   153  }