open-consul/vendor/github.com/hashicorp/go-raftchunking/chunking.go

104 lines
2.7 KiB
Go
Raw Normal View History

package raftchunking
import "github.com/mitchellh/copystructure"
type ChunkStorage interface {
// StoreChunk stores Data from ChunkInfo according to the other metadata
// (OpNum, SeqNum). The bool returns whether or not all chunks have been
// received, as in, the number of non-nil chunks is the same as NumChunks.
StoreChunk(*ChunkInfo) (bool, error)
// FinalizeOp gets all chunks for an op number and then removes the chunk
// info for that op from the store. It should only be called when
// StoreChunk for a given op number returns true but should be safe to call
// at any time; clearing an op can be accomplished by calling this function
// and ignoring the non-error result.
FinalizeOp(uint64) ([]*ChunkInfo, error)
// GetState gets all currently tracked ops, for snapshotting
GetChunks() (ChunkMap, error)
// RestoreChunks restores the current FSM state from a map
RestoreChunks(ChunkMap) error
}
type State struct {
ChunkMap ChunkMap
}
// ChunkInfo holds chunk information
type ChunkInfo struct {
OpNum uint64
SequenceNum uint32
NumChunks uint32
Term uint64
Data []byte
}
// ChunkMap represents a set of data chunks. We use ChunkInfo with Data instead
// of bare []byte in case there is a need to extend this info later.
type ChunkMap map[uint64][]*ChunkInfo
// InmemChunkStorage satisfies ChunkStorage using an in-memory-only tracking
// method.
type InmemChunkStorage struct {
chunks ChunkMap
}
func NewInmemChunkStorage() *InmemChunkStorage {
return &InmemChunkStorage{
chunks: make(ChunkMap),
}
}
func (i *InmemChunkStorage) StoreChunk(chunk *ChunkInfo) (bool, error) {
chunks, ok := i.chunks[chunk.OpNum]
if !ok {
chunks = make([]*ChunkInfo, chunk.NumChunks)
i.chunks[chunk.OpNum] = chunks
}
chunks[chunk.SequenceNum] = chunk
for _, c := range chunks {
// Check for nil, but also check data length in case it ends up
// unmarshaling weirdly for some reason where it makes a new struct
// instead of keeping the pointer nil
if c == nil || len(c.Data) == 0 {
// Not done yet, so return
return false, nil
}
}
return true, nil
}
func (i *InmemChunkStorage) FinalizeOp(opNum uint64) ([]*ChunkInfo, error) {
ret := i.chunks[opNum]
delete(i.chunks, opNum)
return ret, nil
}
func (i *InmemChunkStorage) GetChunks() (ChunkMap, error) {
ret, err := copystructure.Copy(i.chunks)
if err != nil {
return nil, err
}
return ret.(ChunkMap), nil
}
func (i *InmemChunkStorage) RestoreChunks(chunks ChunkMap) error {
// If passed in explicit emptiness, set state to empty
if chunks == nil || len(chunks) == 0 {
i.chunks = make(ChunkMap)
return nil
}
chunksCopy, err := copystructure.Copy(chunks)
if err != nil {
return err
}
i.chunks = chunksCopy.(ChunkMap)
return nil
}