github.com/TBD54566975/ftl@v0.219.0/internal/exec/circularbuffer.go (about)

     1  package exec
     2  
     3  import (
     4  	"bufio"
     5  	"bytes"
     6  	"container/ring"
     7  	"context"
     8  	"fmt"
     9  	"io"
    10  	"runtime"
    11  	"sync"
    12  
    13  	"github.com/TBD54566975/ftl/internal/log"
    14  )
    15  
    16  type CircularBuffer struct {
    17  	r    *ring.Ring
    18  	size int
    19  	mu   sync.Mutex
    20  	cap  int
    21  }
    22  
    23  func NewCircularBuffer(capacity int) *CircularBuffer {
    24  	return &CircularBuffer{
    25  		r:   ring.New(capacity),
    26  		cap: capacity,
    27  	}
    28  }
    29  
    30  // Write accepts a string and stores it in the buffer.
    31  // It expects entire lines and stores each line as a separate entry in the ring buffer.
    32  func (cb *CircularBuffer) Write(p string) error {
    33  	cb.mu.Lock()
    34  	defer cb.mu.Unlock()
    35  
    36  	cb.r.Value = p
    37  	cb.r = cb.r.Next()
    38  
    39  	if cb.size < cb.cap {
    40  		cb.size++
    41  	}
    42  
    43  	return nil
    44  }
    45  
    46  func (cb *CircularBuffer) Bytes() []byte {
    47  	cb.mu.Lock()
    48  	defer cb.mu.Unlock()
    49  
    50  	if cb.size == 0 {
    51  		return []byte{}
    52  	}
    53  
    54  	var buf bytes.Buffer
    55  	start := cb.r.Move(-cb.size) // Correctly calculate the starting position
    56  
    57  	for range cb.size {
    58  		if str, ok := start.Value.(string); ok {
    59  			buf.WriteString(str)
    60  		} else {
    61  			fmt.Println("Unexpected type or nil found in buffer")
    62  		}
    63  		start = start.Next()
    64  	}
    65  
    66  	return buf.Bytes()
    67  }
    68  
    69  func (cb *CircularBuffer) WriterAt(ctx context.Context, level log.Level) *io.PipeWriter {
    70  	// Copied from logger.go which is based on logrus
    71  	// Based on MIT licensed Logrus https://github.com/sirupsen/logrus/blob/bdc0db8ead3853c56b7cd1ac2ba4e11b47d7da6b/writer.go#L27
    72  	logger := log.FromContext(ctx)
    73  	reader, writer := io.Pipe()
    74  	var printFunc func(format string, args ...interface{})
    75  
    76  	switch level {
    77  	case log.Trace:
    78  		printFunc = logger.Tracef
    79  	case log.Debug:
    80  		printFunc = logger.Debugf
    81  	case log.Info:
    82  		printFunc = logger.Infof
    83  	case log.Warn:
    84  		printFunc = logger.Warnf
    85  	case log.Error:
    86  		printFunc = func(format string, args ...interface{}) {
    87  			logger.Errorf(nil, format, args...)
    88  		}
    89  	default:
    90  		panic(level)
    91  	}
    92  
    93  	go cb.writerScanner(reader, printFunc)
    94  	runtime.SetFinalizer(writer, writerFinalizer)
    95  
    96  	return writer
    97  }
    98  
    99  func (cb *CircularBuffer) writerScanner(reader *io.PipeReader, printFunc func(format string, args ...interface{})) {
   100  	scanner := bufio.NewScanner(reader)
   101  	for scanner.Scan() {
   102  		text := scanner.Text()
   103  		printFunc(text)
   104  		err := cb.Write(text + "\n")
   105  		if err != nil {
   106  			fmt.Println("Error writing to buffer")
   107  		}
   108  	}
   109  	if err := scanner.Err(); err != nil {
   110  		fmt.Println("Error reading from pipe:", err)
   111  	}
   112  	reader.Close()
   113  }
   114  
   115  func writerFinalizer(writer *io.PipeWriter) {
   116  	writer.Close()
   117  }