github.com/Finschia/finschia-sdk@v0.48.1/snapshots/stream.go (about)

     1  package snapshots
     2  
     3  import (
     4  	"bufio"
     5  	"compress/zlib"
     6  	"io"
     7  
     8  	protoio "github.com/gogo/protobuf/io"
     9  	"github.com/gogo/protobuf/proto"
    10  
    11  	sdkerrors "github.com/Finschia/finschia-sdk/types/errors"
    12  )
    13  
    14  const (
    15  	// Do not change chunk size without new snapshot format (must be uniform across nodes)
    16  	snapshotChunkSize  = uint64(10e6)
    17  	snapshotBufferSize = int(snapshotChunkSize)
    18  	// Do not change compression level without new snapshot format (must be uniform across nodes)
    19  	snapshotCompressionLevel = 7
    20  )
    21  
    22  // StreamWriter set up a stream pipeline to serialize snapshot nodes:
    23  // Exported Items -> delimited Protobuf -> zlib -> buffer -> chunkWriter -> chan io.ReadCloser
    24  type StreamWriter struct {
    25  	chunkWriter *ChunkWriter
    26  	bufWriter   *bufio.Writer
    27  	zWriter     *zlib.Writer
    28  	protoWriter protoio.WriteCloser
    29  }
    30  
    31  // NewStreamWriter set up a stream pipeline to serialize snapshot DB records.
    32  func NewStreamWriter(ch chan<- io.ReadCloser) *StreamWriter {
    33  	chunkWriter := NewChunkWriter(ch, snapshotChunkSize)
    34  	bufWriter := bufio.NewWriterSize(chunkWriter, snapshotBufferSize)
    35  	zWriter, err := zlib.NewWriterLevel(bufWriter, snapshotCompressionLevel)
    36  	if err != nil {
    37  		chunkWriter.CloseWithError(sdkerrors.Wrap(err, "zlib failure"))
    38  		return nil
    39  	}
    40  	protoWriter := protoio.NewDelimitedWriter(zWriter)
    41  	return &StreamWriter{
    42  		chunkWriter: chunkWriter,
    43  		bufWriter:   bufWriter,
    44  		zWriter:     zWriter,
    45  		protoWriter: protoWriter,
    46  	}
    47  }
    48  
    49  // WriteMsg implements protoio.Write interface
    50  func (sw *StreamWriter) WriteMsg(msg proto.Message) error {
    51  	return sw.protoWriter.WriteMsg(msg)
    52  }
    53  
    54  // Close implements io.Closer interface
    55  func (sw *StreamWriter) Close() error {
    56  	if err := sw.protoWriter.Close(); err != nil {
    57  		sw.chunkWriter.CloseWithError(err)
    58  		return err
    59  	}
    60  	if err := sw.zWriter.Close(); err != nil {
    61  		sw.chunkWriter.CloseWithError(err)
    62  		return err
    63  	}
    64  	if err := sw.bufWriter.Flush(); err != nil {
    65  		sw.chunkWriter.CloseWithError(err)
    66  		return err
    67  	}
    68  	return sw.chunkWriter.Close()
    69  }
    70  
    71  // CloseWithError pass error to chunkWriter
    72  func (sw *StreamWriter) CloseWithError(err error) {
    73  	sw.chunkWriter.CloseWithError(err)
    74  }
    75  
    76  // StreamReader set up a restore stream pipeline
    77  // chan io.ReadCloser -> chunkReader -> zlib -> delimited Protobuf -> ExportNode
    78  type StreamReader struct {
    79  	chunkReader *ChunkReader
    80  	zReader     io.ReadCloser
    81  	protoReader protoio.ReadCloser
    82  }
    83  
    84  // NewStreamReader set up a restore stream pipeline.
    85  func NewStreamReader(chunks <-chan io.ReadCloser) (*StreamReader, error) {
    86  	chunkReader := NewChunkReader(chunks)
    87  	zReader, err := zlib.NewReader(chunkReader)
    88  	if err != nil {
    89  		return nil, sdkerrors.Wrap(err, "zlib failure")
    90  	}
    91  	protoReader := protoio.NewDelimitedReader(zReader, snapshotMaxItemSize)
    92  	return &StreamReader{
    93  		chunkReader: chunkReader,
    94  		zReader:     zReader,
    95  		protoReader: protoReader,
    96  	}, nil
    97  }
    98  
    99  // ReadMsg implements protoio.Reader interface
   100  func (sr *StreamReader) ReadMsg(msg proto.Message) error {
   101  	return sr.protoReader.ReadMsg(msg)
   102  }
   103  
   104  // Close implements io.Closer interface
   105  func (sr *StreamReader) Close() error {
   106  	sr.protoReader.Close()
   107  	sr.zReader.Close()
   108  	return sr.chunkReader.Close()
   109  }