github.com/grailbio/base@v0.0.11/limitbuf/limitbuf.go (about)

     1  package limitbuf
     2  
     3  import (
     4  	"fmt"
     5  	"strings"
     6  
     7  	"github.com/grailbio/base/log"
     8  )
     9  
    10  type (
    11  	// Logger is like strings.Builder, but with maximum length.  If the caller tries
    12  	// to add data beyond the capacity, they will be dropped, and Logger.String()
    13  	// will append "(truncated)" at the end.
    14  	//
    15  	// TODO: Consider renaming to Builder or Buffer since this type's behavior
    16  	// is analogous to those.
    17  	Logger struct {
    18  		maxLen int
    19  		b      strings.Builder
    20  
    21  		// seen counts the total number of bytes passed to Write.
    22  		seen                       int64
    23  		logIfTruncatingMaxMultiple float64
    24  	}
    25  
    26  	// LoggerOption is passed to NewLogger to configure a Logger.
    27  	LoggerOption func(*Logger)
    28  )
    29  
    30  // LogIfTruncatingMaxMultiple controls informative logging about how much data
    31  // passed to Write has been truncated.
    32  //
    33  // If zero, this logging is disabled. Otherwise, if the sum of len(data)
    34  // passed to prior Write calls is greater than LogIfTruncatingMaxMultiple *
    35  // maxLen (passed to NewLogger), a log message is written in the next call
    36  // to String(). After logging, LogIfTruncatingMaxMultiple is set to zero
    37  // to avoid repeating the same message.
    38  //
    39  // This can be a useful diagnostic for both CPU and memory usage if a huge
    40  // amount of data is written and only a tiny fraction is used. For example,
    41  // if a caller writes to the log with fmt.Fprint(logger, ...) they may
    42  // not realize that fmt.Fprint* actually buffers the *entire* formatted
    43  // string in memory first, then writes it to logger.
    44  // TODO: Consider serving the fmt use case better for e.g. bigslice.
    45  //
    46  // Note that the log message is written to log.Error, not the Logger itself
    47  // (it's not part of String's return).
    48  func LogIfTruncatingMaxMultiple(m float64) LoggerOption {
    49  	return func(l *Logger) { l.logIfTruncatingMaxMultiple = m }
    50  }
    51  
    52  // NewLogger creates a new Logger object with the given capacity.
    53  func NewLogger(maxLen int, opts ...LoggerOption) *Logger {
    54  	l := Logger{maxLen: maxLen}
    55  	for _, opt := range opts {
    56  		opt(&l)
    57  	}
    58  	return &l
    59  }
    60  
    61  // Write implements io.Writer interface.
    62  func (b *Logger) Write(data []byte) (int, error) {
    63  	n := b.maxLen - b.b.Len()
    64  	if n > len(data) {
    65  		n = len(data)
    66  	}
    67  	if n > 0 {
    68  		b.b.Write(data[:n])
    69  	}
    70  	b.seen += int64(len(data))
    71  	return len(data), nil
    72  }
    73  
    74  // String reports the data written so far. If the length of the data exceeds the
    75  // buffer capacity, the prefix of the data, plus "(truncated)" will be reported.
    76  func (b *Logger) String() string {
    77  	if b.seen <= int64(b.maxLen) {
    78  		return b.b.String()
    79  	}
    80  	// Truncated.
    81  	if b.logIfTruncatingMaxMultiple > 0 &&
    82  		b.seen > int64(float64(b.maxLen)*b.logIfTruncatingMaxMultiple) {
    83  		b.logIfTruncatingMaxMultiple = 0
    84  		log.Errorf("limitbuf: extreme truncation: %d -> %d bytes", b.seen, b.maxLen)
    85  	}
    86  	return b.b.String() + fmt.Sprintf("(truncated %d bytes)", b.seen-int64(b.maxLen))
    87  }