github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/network/p2p/unicast/protocols/internal/compressedStream.go (about)

     1  package internal
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"sync"
     7  
     8  	"github.com/libp2p/go-libp2p/core/network"
     9  	"go.uber.org/multierr"
    10  
    11  	flownet "github.com/onflow/flow-go/network"
    12  )
    13  
    14  // CompressedStream is an internal networking layer data structure,
    15  // which implements a compression mechanism as a wrapper around a native
    16  // libp2p stream.
    17  type CompressedStream struct {
    18  	network.Stream
    19  
    20  	writeLock  sync.Mutex
    21  	readLock   sync.Mutex
    22  	compressor flownet.Compressor
    23  
    24  	r io.ReadCloser
    25  	w flownet.WriteCloseFlusher
    26  }
    27  
    28  // NewCompressedStream creates a compressed stream with gzip as default compressor.
    29  func NewCompressedStream(s network.Stream, compressor flownet.Compressor) (*CompressedStream, error) {
    30  	c := &CompressedStream{
    31  		Stream:     s,
    32  		compressor: compressor,
    33  	}
    34  
    35  	w, err := c.compressor.NewWriter(s)
    36  	if err != nil {
    37  		return nil, fmt.Errorf("could not create compressor writer: %w", err)
    38  	}
    39  
    40  	c.w = w
    41  
    42  	return c, nil
    43  }
    44  
    45  func (c *CompressedStream) Write(b []byte) (int, error) {
    46  	c.writeLock.Lock()
    47  	defer c.writeLock.Unlock()
    48  
    49  	n, err := c.w.Write(b)
    50  
    51  	return n, multierr.Combine(err, c.w.Flush())
    52  }
    53  
    54  func (c *CompressedStream) Read(b []byte) (int, error) {
    55  	c.readLock.Lock()
    56  	defer c.readLock.Unlock()
    57  
    58  	if c.r == nil {
    59  		r, err := c.compressor.NewReader(c.Stream)
    60  		if err != nil {
    61  			return 0, fmt.Errorf("could not create compressor reader: %w", err)
    62  		}
    63  
    64  		c.r = r
    65  	}
    66  
    67  	n, err := c.r.Read(b)
    68  	if err != nil {
    69  		c.r.Close()
    70  	}
    71  	return n, err
    72  }
    73  
    74  func (c *CompressedStream) Close() error {
    75  	c.writeLock.Lock()
    76  	defer c.writeLock.Unlock()
    77  
    78  	return multierr.Combine(c.w.Close(), c.Stream.Close())
    79  }