github.com/haraldrudell/parl@v0.4.176/pio/tap.go (about)

     1  /*
     2  © 2023–present Harald Rudell <harald.rudell@gmail.com> (https://haraldrudell.github.io/haraldrudell/)
     3  ISC License
     4  */
     5  
     6  package pio
     7  
     8  import (
     9  	"io"
    10  	"sync/atomic"
    11  
    12  	"github.com/haraldrudell/parl"
    13  )
    14  
    15  // Tap returns a socket tap producing two streams of data
    16  // read from and written to a socket
    17  type Tap struct {
    18  	closeWinner               atomic.Bool
    19  	IsClosed                  parl.Awaitable
    20  	readsWriter, writesWriter io.Writer
    21  	addError                  func(err error)
    22  }
    23  
    24  func NewTap(readsWriter, writesWriter io.Writer, addError parl.AddError) (tap *Tap) {
    25  	return &Tap{
    26  		readsWriter:  readsWriter,
    27  		writesWriter: writesWriter,
    28  		addError:     addError,
    29  	}
    30  }
    31  
    32  func (t *Tap) Read(reader io.Reader, p []byte) (n int, err error) {
    33  
    34  	// do delegated Read
    35  	n, err = reader.Read(p)
    36  
    37  	// copy read data to reads
    38  	if n > 0 && t.readsWriter != nil {
    39  		var readsN, readsErr = t.readsWriter.Write(p[:n])
    40  		// error in reads
    41  		if readsErr != nil {
    42  			t.handleError(NewPioError(PeReads, readsErr))
    43  		} else if readsN < n {
    44  			// reads short write
    45  			t.handleError(NewPioError(PeReads, io.ErrShortWrite))
    46  		}
    47  	}
    48  
    49  	// propagate reader error
    50  	if err != nil && t.addError != nil {
    51  		t.addError(NewPioError(PeRead, err))
    52  	}
    53  
    54  	return
    55  }
    56  
    57  func (t *Tap) Write(writer io.Writer, p []byte) (n int, err error) {
    58  
    59  	// copy data to writes writer
    60  	if t.writesWriter != nil {
    61  		var writesN, writesErr = t.writesWriter.Write(p)
    62  		// error in writes
    63  		if writesErr != nil {
    64  			t.handleError(NewPioError(PeWrites, writesErr))
    65  		} else if writesN < n {
    66  			// reads short write
    67  			t.handleError(NewPioError(PeWrites, io.ErrShortWrite))
    68  		}
    69  	}
    70  
    71  	// do delegated Write
    72  	n, err = writer.Write(p)
    73  
    74  	// propagate reader error to addError as well
    75  	if err != nil && t.addError != nil {
    76  		t.addError(NewPioError(PeWrite, err))
    77  	}
    78  
    79  	return
    80  }
    81  
    82  func (t *Tap) Close(closer any) (err error) {
    83  
    84  	// pick closing invocation
    85  	if !t.closeWinner.CompareAndSwap(false, true) {
    86  		<-t.IsClosed.Ch()
    87  		return
    88  	}
    89  	defer t.IsClosed.Close()
    90  
    91  	// close delegate if it implements io.Close
    92  	if closer, ok := closer.(io.Closer); ok {
    93  		if parl.Close(closer, &err); err != nil && t.addError != nil {
    94  			t.addError(NewPioError(PeClose, err))
    95  		}
    96  	}
    97  
    98  	// reads and writes
    99  	var e [2]error
   100  	for i, a := range []any{t.readsWriter, t.writesWriter} {
   101  		var closer, ok = a.(io.Closer)
   102  		if !ok {
   103  			continue
   104  		}
   105  		parl.Close(closer, &e[i])
   106  	}
   107  
   108  	// handle errors, may panic
   109  	for i, source := range []PIOErrorSource{PeReads, PeWrites} {
   110  		if e[i] == nil {
   111  			continue
   112  		}
   113  		t.handleError(NewPioError(source, e[i]))
   114  	}
   115  
   116  	return
   117  }
   118  
   119  func (t *Tap) handleError(err error) {
   120  	if t.addError != nil {
   121  		t.addError(err)
   122  	} else {
   123  		panic(err)
   124  	}
   125  }
   126  
   127  // MultiWriter creates a writer that duplicates its writes to all the provided writers
   128  //   - func io.MultiWriter(writers ...io.Writer) io.Writer
   129  var _ = io.MultiWriter
   130  
   131  // TeeReader returns a Reader that writes to w what it reads from r
   132  //   - func io.TeeReader(r io.Reader, w io.Writer) io.Reader
   133  var _ = io.TeeReader
   134  
   135  // Writer is the interface that wraps the basic Write method
   136  var _ io.Writer
   137  
   138  // ReadWriter is the interface that groups the basic Read and Write methods.
   139  var _ io.ReadWriter