github.com/badrootd/celestia-core@v0.0.0-20240305091328-aa4207a4b25d/test/e2e/app/snapshots.go (about)

     1  package app
     2  
     3  import (
     4  	"encoding/json"
     5  	"errors"
     6  	"fmt"
     7  	"math"
     8  	"os"
     9  	"path/filepath"
    10  	"sync"
    11  
    12  	abci "github.com/badrootd/celestia-core/abci/types"
    13  )
    14  
    15  const (
    16  	snapshotChunkSize = 1e6
    17  )
    18  
    19  // SnapshotStore stores state sync snapshots. Snapshots are stored simply as
    20  // JSON files, and chunks are generated on-the-fly by splitting the JSON data
    21  // into fixed-size chunks.
    22  type SnapshotStore struct {
    23  	sync.RWMutex
    24  	dir      string
    25  	metadata []abci.Snapshot
    26  }
    27  
    28  // NewSnapshotStore creates a new snapshot store.
    29  func NewSnapshotStore(dir string) (*SnapshotStore, error) {
    30  	store := &SnapshotStore{dir: dir}
    31  	if err := os.MkdirAll(dir, 0o755); err != nil {
    32  		return nil, err
    33  	}
    34  	if err := store.loadMetadata(); err != nil {
    35  		return nil, err
    36  	}
    37  	return store, nil
    38  }
    39  
    40  // loadMetadata loads snapshot metadata. Does not take out locks, since it's
    41  // called internally on construction.
    42  func (s *SnapshotStore) loadMetadata() error {
    43  	file := filepath.Join(s.dir, "metadata.json")
    44  	metadata := []abci.Snapshot{}
    45  
    46  	bz, err := os.ReadFile(file)
    47  	switch {
    48  	case errors.Is(err, os.ErrNotExist):
    49  	case err != nil:
    50  		return fmt.Errorf("failed to load snapshot metadata from %q: %w", file, err)
    51  	}
    52  	if len(bz) != 0 {
    53  		err = json.Unmarshal(bz, &metadata)
    54  		if err != nil {
    55  			return fmt.Errorf("invalid snapshot data in %q: %w", file, err)
    56  		}
    57  	}
    58  	s.metadata = metadata
    59  	return nil
    60  }
    61  
    62  // saveMetadata saves snapshot metadata. Does not take out locks, since it's
    63  // called internally from e.g. Create().
    64  func (s *SnapshotStore) saveMetadata() error {
    65  	bz, err := json.Marshal(s.metadata)
    66  	if err != nil {
    67  		return err
    68  	}
    69  
    70  	// save the file to a new file and move it to make saving atomic.
    71  	newFile := filepath.Join(s.dir, "metadata.json.new")
    72  	file := filepath.Join(s.dir, "metadata.json")
    73  	err = os.WriteFile(newFile, bz, 0o644) //nolint: gosec
    74  	if err != nil {
    75  		return err
    76  	}
    77  	return os.Rename(newFile, file)
    78  }
    79  
    80  // Create creates a snapshot of the given application state's key/value pairs.
    81  func (s *SnapshotStore) Create(state *State) (abci.Snapshot, error) {
    82  	s.Lock()
    83  	defer s.Unlock()
    84  	bz, err := state.Export()
    85  	if err != nil {
    86  		return abci.Snapshot{}, err
    87  	}
    88  	snapshot := abci.Snapshot{
    89  		Height: state.Height,
    90  		Format: 1,
    91  		Hash:   hashItems(state.Values),
    92  		Chunks: byteChunks(bz),
    93  	}
    94  	//nolint:gosec // G306: Expect WriteFile permissions to be 0600 or less
    95  	err = os.WriteFile(filepath.Join(s.dir, fmt.Sprintf("%v.json", state.Height)), bz, 0o644)
    96  	if err != nil {
    97  		return abci.Snapshot{}, err
    98  	}
    99  	s.metadata = append(s.metadata, snapshot)
   100  	err = s.saveMetadata()
   101  	if err != nil {
   102  		return abci.Snapshot{}, err
   103  	}
   104  	return snapshot, nil
   105  }
   106  
   107  // List lists available snapshots.
   108  func (s *SnapshotStore) List() ([]*abci.Snapshot, error) {
   109  	s.RLock()
   110  	defer s.RUnlock()
   111  	snapshots := make([]*abci.Snapshot, len(s.metadata))
   112  	for idx := range s.metadata {
   113  		snapshots[idx] = &s.metadata[idx]
   114  	}
   115  	return snapshots, nil
   116  }
   117  
   118  // LoadChunk loads a snapshot chunk.
   119  func (s *SnapshotStore) LoadChunk(height uint64, format uint32, chunk uint32) ([]byte, error) {
   120  	s.RLock()
   121  	defer s.RUnlock()
   122  	for _, snapshot := range s.metadata {
   123  		if snapshot.Height == height && snapshot.Format == format {
   124  			bz, err := os.ReadFile(filepath.Join(s.dir, fmt.Sprintf("%v.json", height)))
   125  			if err != nil {
   126  				return nil, err
   127  			}
   128  			return byteChunk(bz, chunk), nil
   129  		}
   130  	}
   131  	return nil, nil
   132  }
   133  
   134  // byteChunk returns the chunk at a given index from the full byte slice.
   135  func byteChunk(bz []byte, index uint32) []byte {
   136  	start := int(index * snapshotChunkSize)
   137  	end := int((index + 1) * snapshotChunkSize)
   138  	switch {
   139  	case start >= len(bz):
   140  		return nil
   141  	case end >= len(bz):
   142  		return bz[start:]
   143  	default:
   144  		return bz[start:end]
   145  	}
   146  }
   147  
   148  // byteChunks calculates the number of chunks in the byte slice.
   149  func byteChunks(bz []byte) uint32 {
   150  	return uint32(math.Ceil(float64(len(bz)) / snapshotChunkSize))
   151  }