github.com/searKing/golang/go@v1.2.117/io/sniff.reader.go (about)

     1  // Copyright 2021 The searKing Author. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package io
     6  
     7  import (
     8  	"bytes"
     9  	"io"
    10  )
    11  
    12  type sniffReader struct {
    13  	// sniff start: read from [historyBuffers..., source] and buffered in buffer
    14  	// sniff stop: read from [buffer, historyBuffers..., source] and clean buffer and historyBuffers if meet EOF
    15  	source io.Reader
    16  
    17  	// virtual reader: buffer, historyBuffers..., source
    18  	buffer         *bytes.Buffer // latest read data
    19  	historyBuffers []io.Reader   // new, old, older read data
    20  
    21  	selectorF DynamicReaderFunc
    22  
    23  	sniffing bool
    24  }
    25  
    26  func newSniffReader(r io.Reader) *sniffReader {
    27  	sr := &sniffReader{
    28  		source: r,
    29  	}
    30  	sr.stopSniff()
    31  	return sr
    32  }
    33  
    34  // Sniff starts or stops sniffing, restarts if stop and start called one by one
    35  // true to start sniffing all data unread actually
    36  // false to return a multi reader with all data sniff buffered and source
    37  func (sr *sniffReader) Sniff(sniffing bool) ReadSniffer {
    38  	if sr.sniffing == sniffing {
    39  		return sr
    40  	}
    41  	sr.sniffing = sniffing
    42  	if sniffing {
    43  		sr.startSniff()
    44  		return sr
    45  	}
    46  	sr.stopSniff()
    47  	return sr
    48  }
    49  
    50  // shrinkToHistory shrink buffer to history buffers
    51  func (sr *sniffReader) shrinkToHistory() {
    52  	if sr.buffer != nil {
    53  		if sr.buffer.Len() > 0 {
    54  			// clear if EOF meet
    55  			bufferReader := WatchReader(bytes.NewBuffer(sr.buffer.Bytes()), WatcherFunc(func(p []byte, n int, err error) (int, error) {
    56  				if err == io.EOF {
    57  					// historyBuffers is consumed head first, so can be cleared from head
    58  					sr.historyBuffers = sr.historyBuffers[1:] // remove head to recover space
    59  				}
    60  				return n, err
    61  			}))
    62  			sr.historyBuffers = append([]io.Reader{bufferReader}, sr.historyBuffers...)
    63  		}
    64  		sr.buffer = nil
    65  	}
    66  }
    67  
    68  // startSniff starts sniff and return a TeeReader that writes to buffer while reads from history buffers and source
    69  func (sr *sniffReader) startSniff() {
    70  	sr.shrinkToHistory()
    71  	// We don't need the buffer anymore.
    72  	// Reset it to release the internal slice.
    73  	sr.buffer = &bytes.Buffer{}
    74  
    75  	readers := append(sr.historyBuffers, sr.source)
    76  	reader := io.TeeReader(io.MultiReader(readers...), sr.buffer)
    77  	sr.selectorF = func() io.Reader { return reader }
    78  }
    79  
    80  // stopSniff stops sniff and return a MultiReader of history buffers and source
    81  func (sr *sniffReader) stopSniff() {
    82  	sr.shrinkToHistory()
    83  	readers := append(sr.historyBuffers, sr.source)
    84  	reader := io.MultiReader(readers...)
    85  	sr.selectorF = func() io.Reader { return reader }
    86  }
    87  
    88  func (sr *sniffReader) Read(p []byte) (n int, err error) { return sr.selectorF.Read(p) }