git.sr.ht/~pingoo/stdx@v0.0.0-20240218134121-094174641f6e/log/loki/writer.go (about)

     1  package loki
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"io"
     7  	"log"
     8  	"net/http"
     9  	"os"
    10  	"sync"
    11  	"time"
    12  
    13  	"git.sr.ht/~pingoo/stdx/httpx"
    14  )
    15  
    16  type WriterOptions struct {
    17  	LokiEndpoint string
    18  	// ChildWriter is used to pass logs to another writer
    19  	// default: os.Stdout
    20  	ChildWriter io.Writer
    21  	// DefaultRecordsBufferSize is your expected number of `(logs per second) / (1000 / FlushTimeout)`
    22  	// default: 100
    23  	DefaultRecordsBufferSize uint32
    24  	// EmptyEndpointMaxBufferSize is the number of logs to buffer if LokiEndpoint == "".
    25  	// it's useful if you log a few things before your config with the loki endpoint is full loaded
    26  	// default: 200
    27  	EmptyEndpointMaxBufferSize uint32
    28  	// FlushTimeout in ms
    29  	// default: 200
    30  	FlushTimeout uint32
    31  }
    32  
    33  type Writer struct {
    34  	lokiEndpoint               string
    35  	streams                    map[string]string
    36  	defaultRecordsBufferSize   uint32
    37  	emptyEndpointMaxBufferSize uint32
    38  	flushTimeout               uint32
    39  
    40  	httpClient         *http.Client
    41  	recordsBuffer      []record
    42  	recordsBufferMutex sync.Mutex
    43  	childWriter        io.Writer
    44  	ctx                context.Context
    45  }
    46  
    47  type record struct {
    48  	timestamp time.Time
    49  	message   string
    50  }
    51  
    52  func NewWriter(ctx context.Context, lokiEndpoint string, streams map[string]string, options *WriterOptions) *Writer {
    53  	if streams == nil {
    54  		streams = map[string]string{}
    55  	}
    56  
    57  	defaultOptions := defaultOptions()
    58  	if options == nil {
    59  		options = defaultOptions
    60  	} else {
    61  		if options.ChildWriter == nil {
    62  			options.ChildWriter = defaultOptions.ChildWriter
    63  		}
    64  		if options.DefaultRecordsBufferSize == 0 {
    65  			options.DefaultRecordsBufferSize = defaultOptions.DefaultRecordsBufferSize
    66  		}
    67  		if options.EmptyEndpointMaxBufferSize == 0 {
    68  			options.EmptyEndpointMaxBufferSize = defaultOptions.EmptyEndpointMaxBufferSize
    69  		}
    70  		if options.FlushTimeout == 0 {
    71  			options.FlushTimeout = defaultOptions.FlushTimeout
    72  		}
    73  	}
    74  
    75  	if ctx == nil {
    76  		ctx = context.Background()
    77  	}
    78  
    79  	writer := &Writer{
    80  		lokiEndpoint:               lokiEndpoint,
    81  		streams:                    streams,
    82  		defaultRecordsBufferSize:   options.DefaultRecordsBufferSize,
    83  		emptyEndpointMaxBufferSize: options.EmptyEndpointMaxBufferSize,
    84  		flushTimeout:               options.FlushTimeout,
    85  
    86  		httpClient:         httpx.DefaultClient(),
    87  		recordsBuffer:      make([]record, 0, options.DefaultRecordsBufferSize),
    88  		recordsBufferMutex: sync.Mutex{},
    89  		childWriter:        options.ChildWriter,
    90  		ctx:                ctx,
    91  	}
    92  
    93  	go func() {
    94  		done := false
    95  		for {
    96  			if done {
    97  				// we sleep less to avoid losing logs
    98  				time.Sleep(20 * time.Millisecond)
    99  			} else {
   100  				select {
   101  				case <-writer.ctx.Done():
   102  					done = true
   103  				case <-time.After(time.Duration(writer.flushTimeout) * time.Millisecond):
   104  				}
   105  			}
   106  
   107  			go func() {
   108  				// TODO: as of now, if the HTTP request fail after X retries, we discard/lose the logs
   109  				err := writer.flushLogs(context.Background())
   110  				if err != nil {
   111  					log.Println(err.Error())
   112  					return
   113  				}
   114  				// if err != nil {
   115  				// 		writer.recordsBufferMutex.Lock()
   116  				// 		writer.recordsBuffer = append(writer.recordsBuffer, recordsBufferCopy...)
   117  				// 		writer.recordsBufferMutex.Unlock()
   118  				// 	}
   119  			}()
   120  		}
   121  	}()
   122  
   123  	return writer
   124  }
   125  
   126  func defaultOptions() *WriterOptions {
   127  	return &WriterOptions{
   128  		LokiEndpoint:               "",
   129  		ChildWriter:                os.Stdout,
   130  		DefaultRecordsBufferSize:   100,
   131  		EmptyEndpointMaxBufferSize: 200,
   132  		FlushTimeout:               200,
   133  	}
   134  }
   135  
   136  // SetEndpoint sets the loki endpoint. This method IS NOT thread safe.
   137  // It should be used just after config is loaded
   138  func (writer *Writer) SetEndpoint(lokiEndpoint string) {
   139  	writer.lokiEndpoint = lokiEndpoint
   140  }
   141  
   142  func (writer *Writer) Write(data []byte) (n int, err error) {
   143  	// TODO: handle error?
   144  	_, _ = writer.childWriter.Write(data)
   145  
   146  	// if log finishes by '\n' we trim it
   147  	data = bytes.TrimSuffix(data, []byte("\n"))
   148  
   149  	record := record{
   150  		timestamp: time.Now().UTC(),
   151  		message:   string(data),
   152  	}
   153  
   154  	writer.recordsBufferMutex.Lock()
   155  	if writer.lokiEndpoint != "" || len(writer.recordsBuffer) < int(writer.emptyEndpointMaxBufferSize) {
   156  		writer.recordsBuffer = append(writer.recordsBuffer, record)
   157  	}
   158  	writer.recordsBufferMutex.Unlock()
   159  
   160  	return
   161  }