github.com/go4org/go4@v0.0.0-20200104003542-c7e774b10ea0/writerutil/writerutil.go (about)

     1  /*
     2  Copyright 2016 The go4 Authors
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8       http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  // Package writerutil contains io.Writer types.
    18  package writerutil // import "go4.org/writerutil"
    19  
    20  import (
    21  	"bytes"
    22  	"strconv"
    23  )
    24  
    25  // PrefixSuffixSaver is an io.Writer which retains the first N bytes
    26  // and the last N bytes written to it. The Bytes method reconstructs
    27  // it with a pretty error message.
    28  // It is copied from os/exec/exec.go of the Go stdlib.
    29  type PrefixSuffixSaver struct {
    30  	N         int // max size of prefix or suffix
    31  	prefix    []byte
    32  	suffix    []byte // ring buffer once len(suffix) == N
    33  	suffixOff int    // offset to write into suffix
    34  	skipped   int64
    35  
    36  	// TODO(bradfitz): we could keep one large []byte and use part of it for
    37  	// the prefix, reserve space for the '... Omitting N bytes ...' message,
    38  	// then the ring buffer suffix, and just rearrange the ring buffer
    39  	// suffix when Bytes() is called, but it doesn't seem worth it for
    40  	// now just for error messages. It's only ~64KB anyway.
    41  }
    42  
    43  func (w *PrefixSuffixSaver) Write(p []byte) (n int, err error) {
    44  	lenp := len(p)
    45  	p = w.fill(&w.prefix, p)
    46  
    47  	// Only keep the last w.N bytes of suffix data.
    48  	if overage := len(p) - w.N; overage > 0 {
    49  		p = p[overage:]
    50  		w.skipped += int64(overage)
    51  	}
    52  	p = w.fill(&w.suffix, p)
    53  
    54  	// w.suffix is full now if p is non-empty. Overwrite it in a circle.
    55  	for len(p) > 0 { // 0, 1, or 2 iterations.
    56  		n := copy(w.suffix[w.suffixOff:], p)
    57  		p = p[n:]
    58  		w.skipped += int64(n)
    59  		w.suffixOff += n
    60  		if w.suffixOff == w.N {
    61  			w.suffixOff = 0
    62  		}
    63  	}
    64  	return lenp, nil
    65  }
    66  
    67  // fill appends up to len(p) bytes of p to *dst, such that *dst does not
    68  // grow larger than w.N. It returns the un-appended suffix of p.
    69  func (w *PrefixSuffixSaver) fill(dst *[]byte, p []byte) (pRemain []byte) {
    70  	if remain := w.N - len(*dst); remain > 0 {
    71  		add := minInt(len(p), remain)
    72  		*dst = append(*dst, p[:add]...)
    73  		p = p[add:]
    74  	}
    75  	return p
    76  }
    77  
    78  // Bytes returns a slice of the bytes, or a copy of the bytes, retained by w.
    79  // If more bytes than could be retained were written to w, it returns a
    80  // concatenation of the N first bytes, a message for how many bytes were dropped,
    81  // and the N last bytes.
    82  func (w *PrefixSuffixSaver) Bytes() []byte {
    83  	if w.suffix == nil {
    84  		return w.prefix
    85  	}
    86  	if w.skipped == 0 {
    87  		return append(w.prefix, w.suffix...)
    88  	}
    89  	var buf bytes.Buffer
    90  	buf.Grow(len(w.prefix) + len(w.suffix) + 50)
    91  	buf.Write(w.prefix)
    92  	buf.WriteString("\n... omitting ")
    93  	buf.WriteString(strconv.FormatInt(w.skipped, 10))
    94  	buf.WriteString(" bytes ...\n")
    95  	buf.Write(w.suffix[w.suffixOff:])
    96  	buf.Write(w.suffix[:w.suffixOff])
    97  	return buf.Bytes()
    98  }
    99  
   100  func minInt(a, b int) int {
   101  	if a < b {
   102  		return a
   103  	}
   104  	return b
   105  }