github.com/mailru/activerecord@v1.12.2/pkg/iproto/util/io/io.go (about)

     1  // Package io contains utility for working with golang's io objects.
     2  package io
     3  
     4  import (
     5  	"io"
     6  	"time"
     7  )
     8  
     9  // Stat represents statistics about underlying io.{Reader,Writer} usage.
    10  type Stat struct {
    11  	Bytes uint32 // Bytes sent/read from underlying object.
    12  	Calls uint32 // Read/Write calls made to the underlying object.
    13  }
    14  
    15  // reader is a wrapper around io.Reader.
    16  // Underlying reader should not be *bufio.Reader.
    17  // It used to calculate stats of reading from underlying reader.
    18  type Reader struct {
    19  	r     io.Reader
    20  	bytes uint32 // bytes read
    21  	calls uint32 // calls made
    22  }
    23  
    24  // WrapReader wraps r into Reader to calculate usage stats of r.
    25  // Note that Reader is not goroutine safe.
    26  func WrapReader(r io.Reader) *Reader {
    27  	ret := &Reader{
    28  		r: r,
    29  	}
    30  
    31  	return ret
    32  }
    33  
    34  // Read implements io.Reader interface.
    35  func (r *Reader) Read(p []byte) (int, error) {
    36  	n, err := r.r.Read(p)
    37  	r.bytes += uint32(n)
    38  	r.calls++
    39  
    40  	return n, err
    41  }
    42  
    43  // Stat returns underlying io.Reader usage statistics.
    44  func (r *Reader) Stat() Stat {
    45  	return Stat{
    46  		Bytes: r.bytes,
    47  		Calls: r.calls,
    48  	}
    49  }
    50  
    51  // Writer is a wrapper around io.Writer.
    52  // Underlying writer should not be *bufio.Writer.
    53  // It used to calculate stats of writing to underlying writer.
    54  type Writer struct {
    55  	w     io.Writer
    56  	bytes uint32 // bytes written
    57  	calls uint32 // calls made
    58  }
    59  
    60  // WrapWriter wraps w into Writer to calculate usage stats of w.
    61  // Note that Writer is not goroutine safe.
    62  func WrapWriter(w io.Writer) *Writer {
    63  	ret := &Writer{
    64  		w: w,
    65  	}
    66  
    67  	return ret
    68  }
    69  
    70  // Write implements io.Writer.
    71  func (w *Writer) Write(p []byte) (int, error) {
    72  	n, err := w.w.Write(p)
    73  	w.bytes += uint32(n)
    74  	w.calls++
    75  
    76  	return n, err
    77  }
    78  
    79  // Stat returns underlying io.Writer usage statistics.
    80  func (w *Writer) Stat() Stat {
    81  	return Stat{
    82  		Bytes: w.bytes,
    83  		Calls: w.calls,
    84  	}
    85  }
    86  
    87  // DeadlineWriter describes object that could prepare io.Writer methods with
    88  // some deadline.
    89  type DeadlineWriter interface {
    90  	io.Writer
    91  	SetWriteDeadline(time.Time) error
    92  }
    93  
    94  // DeadlineWriter describes object that could prepare io.Reader methods with
    95  // some deadline.
    96  type DeadlineReader interface {
    97  	io.Reader
    98  	SetReadDeadline(time.Time) error
    99  }
   100  
   101  // TimeoutWriter is a wrapper around DeadlineWriter that sets write deadline on
   102  // each Write() call. It is useful as destination for bufio.Writer, when you do
   103  // not exactly know, when Write() will occure, but want to control timeout of
   104  // such calls.
   105  type TimeoutWriter struct {
   106  	Dest    DeadlineWriter
   107  	Timeout time.Duration
   108  }
   109  
   110  // Write implements io.Writer interface.
   111  func (w TimeoutWriter) Write(p []byte) (int, error) {
   112  	if err := w.Dest.SetWriteDeadline(time.Now().Add(w.Timeout)); err != nil {
   113  		return 0, err
   114  	}
   115  
   116  	return w.Dest.Write(p)
   117  }
   118  
   119  // TimeoutReader is a wrapper around DeadlineReader that sets read deadline on
   120  // each Read() call. It is useful as destination for bufio.Reader, when you do
   121  // not exactly know, when Read() will occure, but want to control timeout of
   122  // such calls.
   123  type TimeoutReader struct {
   124  	Dest    DeadlineReader
   125  	Timeout time.Duration
   126  }
   127  
   128  // Read implements io.Reader interface.
   129  func (w TimeoutReader) Read(p []byte) (int, error) {
   130  	if err := w.Dest.SetReadDeadline(time.Now().Add(w.Timeout)); err != nil {
   131  		return 0, err
   132  	}
   133  
   134  	return w.Dest.Read(p)
   135  }