go_study/fabric-main/common/ledger/snapshot/file.go

/*
Copyright IBM Corp. All Rights Reserved.

SPDX-License-Identifier: Apache-2.0
*/

package snapshot

import (
	"bufio"
	"encoding/binary"
	"fmt"
	"hash"
	"io"
	"os"

	"github.com/golang/protobuf/proto"
	"github.com/pkg/errors"
)

type NewHashFunc func() (hash.Hash, error)

// FileWriter creates a new file for ledger snapshot. This is expected to be used by various
// components of ledger, such as blockstorage and statedb for exporting the relevant snapshot data
type FileWriter struct {
	file              *os.File
	hasher            hash.Hash
	bufWriter         *bufio.Writer
	multiWriter       io.Writer
	varintReusableBuf []byte
}

// CreateFile creates a new file for exporting the ledger snapshot data
// This function returns an error if the file already exists. The `dataformat` is the first byte
// written to the file. The function newHash is used to construct an hash.Hash for computing the hash-sum of the data stream
func CreateFile(filePath string, dataformat byte, newHashFunc NewHashFunc) (*FileWriter, error) {
	hashImpl, err := newHashFunc()
	if err != nil {
		return nil, err
	}
	// create the file only if it does not already exist.
	// set the permission mode to read-only, as once the file is closed, we do not support modifying the file
	file, err := os.OpenFile(filePath, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0o444)
	if err != nil {
		return nil, errors.Wrapf(err, "error while creating the snapshot file: %s", filePath)
	}
	bufWriter := bufio.NewWriter(file)
	multiWriter := io.MultiWriter(bufWriter, hashImpl)
	if _, err := multiWriter.Write([]byte{dataformat}); err != nil {
		file.Close()
		return nil, errors.Wrapf(err, "error while writing data format to the snapshot file: %s", filePath)
	}
	return &FileWriter{
		file:              file,
		bufWriter:         bufWriter,
		multiWriter:       multiWriter,
		hasher:            hashImpl,
		varintReusableBuf: make([]byte, binary.MaxVarintLen64),
	}, nil
}

// EncodeString encodes and appends the string to the data stream
func (c *FileWriter) EncodeString(str string) error {
	return c.EncodeBytes([]byte(str))
}

// EncodeString encodes and appends a proto message to the data stream
func (c *FileWriter) EncodeProtoMessage(m proto.Message) error {
	b, err := proto.Marshal(m)
	if err != nil {
		return errors.Wrapf(err, "error marshalling proto message to write to the snapshot file: %s", c.file.Name())
	}
	return c.EncodeBytes(b)
}

// EncodeBytes encodes and appends bytes to the data stream
func (c *FileWriter) EncodeBytes(b []byte) error {
	if err := c.EncodeUVarint(uint64(len(b))); err != nil {
		return err
	}
	if _, err := c.multiWriter.Write(b); err != nil {
		return errors.Wrapf(err, "error while writing data to the snapshot file: %s", c.file.Name())
	}
	return nil
}

// EncodeUVarint encodes and appends a number to the data stream
func (c *FileWriter) EncodeUVarint(u uint64) error {
	n := binary.PutUvarint(c.varintReusableBuf, u)
	if _, err := c.multiWriter.Write(c.varintReusableBuf[:n]); err != nil {
		return errors.Wrapf(err, "error while writing data to the snapshot file: %s", c.file.Name())
	}
	return nil
}

// Done closes the snapshot file and returns the final hash of the data stream
func (c *FileWriter) Done() ([]byte, error) {
	if err := c.bufWriter.Flush(); err != nil {
		return nil, errors.Wrapf(err, "error while flushing to the snapshot file: %s ", c.file.Name())
	}
	if err := c.file.Sync(); err != nil {
		return nil, err
	}
	if err := c.file.Close(); err != nil {
		return nil, errors.Wrapf(err, "error while closing the snapshot file: %s ", c.file.Name())
	}
	return c.hasher.Sum(nil), nil
}

// Close closes the underlying file, if not already done. A consumer can invoke this function if the consumer
// encountered some error and simply wants to abandon the snapshot file creation (typically, intended to be used in a defer statement)
func (c *FileWriter) Close() error {
	if c == nil {
		return nil
	}
	return errors.Wrapf(c.file.Close(), "error while closing the snapshot file: %s", c.file.Name())
}

// FileReader reads from a ledger snapshot file. This is expected to be used for loading the ledger snapshot data
// during bootstrapping a channel from snapshot. The data should be read, using the functions `DecodeXXX`,
// in the same sequence in which the data was written by the functions `EncodeXXX` in the `FileCreator`.
// Note that the FileReader does not verify the hash of stream and it is expected that the hash has been verified
// by the consumer. Later, if we decide to perform this, on-the-side, while loading the snapshot data, the FileRedear,
// like the FileCreator, would take a `hasher` as an input
type FileReader struct {
	file              *os.File
	bufReader         *bufio.Reader
	reusableByteSlice []byte
}

// OpenFile constructs a FileReader. This function returns an error if the format of the file, stored in the
// first byte, does not match with the expectedDataFormat
func OpenFile(filePath string, expectDataformat byte) (*FileReader, error) {
	file, err := os.Open(filePath)
	if err != nil {
		return nil, errors.Wrapf(err, "error while opening the snapshot file: %s", filePath)
	}
	bufReader := bufio.NewReader(file)
	dataFormat, err := bufReader.ReadByte()
	if err != nil {
		file.Close()
		return nil, errors.Wrapf(err, "error while reading from the snapshot file: %s", filePath)
	}
	if dataFormat != expectDataformat {
		file.Close()
		return nil, errors.New(fmt.Sprintf("unexpected data format: %x", dataFormat))
	}
	return &FileReader{
		file:      file,
		bufReader: bufReader,
	}, nil
}

// DecodeString reads and decodes a string
func (r *FileReader) DecodeString() (string, error) {
	b, err := r.decodeBytes()
	return string(b), err
}

// DecodeBytes reads and decodes bytes
func (r *FileReader) DecodeBytes() ([]byte, error) {
	b, err := r.decodeBytes()
	if err != nil {
		return nil, err
	}
	c := make([]byte, len(b))
	copy(c, b)
	return c, nil
}

// DecodeUVarInt reads and decodes a number
func (r *FileReader) DecodeUVarInt() (uint64, error) {
	u, err := binary.ReadUvarint(r.bufReader)
	if err != nil {
		return 0, errors.Wrapf(err, "error while reading from snapshot file: %s", r.file.Name())
	}
	return u, nil
}

// DecodeProtoMessage reads and decodes a protoMessage
func (r *FileReader) DecodeProtoMessage(m proto.Message) error {
	b, err := r.decodeBytes()
	if err != nil {
		return err
	}
	return proto.Unmarshal(b, m)
}

// Close closes the file
func (r *FileReader) Close() error {
	if r == nil {
		return nil
	}
	return errors.Wrapf(r.file.Close(), "error while closing the snapshot file: %s", r.file.Name())
}

func (r *FileReader) decodeBytes() ([]byte, error) {
	sizeUint, err := r.DecodeUVarInt()
	if err != nil {
		return nil, err
	}
	size := int(sizeUint)
	if size == 0 {
		return []byte{}, nil
	}
	if len(r.reusableByteSlice) < size {
		r.reusableByteSlice = make([]byte, size)
	}
	if _, err := io.ReadFull(r.bufReader, r.reusableByteSlice[0:size]); err != nil {
		return nil, errors.Wrapf(err, "error while reading from snapshot file: %s", r.file.Name())
	}
	return r.reusableByteSlice[0:size], nil
}