e2e: Add thread-safe getters to app state (celestiaorg#713)

* e2e: Add thread-safe getters to app state

It appears as though a race condition has been triggered in the E2E app
itself due to directly accessing application state inner variables.

Make the application state variables private to remove the temptation to
access them directly from outside of the package and to highlight their
use, and provide thread-safe getters to access the application state
properties.

Signed-off-by: Thane Thomson <[email protected]>

* e2e: Use serializedState to unmarshal state during rollback

Signed-off-by: Thane Thomson <[email protected]>

* Remove unnecessary comment

Signed-off-by: Thane Thomson <[email protected]>

* e2e: Condense thread-safe accessors

Refactor thread-safe accessors such that we do not unlock and re-lock
multiple times during a single ABCI call.

Signed-off-by: Thane Thomson <[email protected]>

* e2e: Use custom JSON marshalling/unmarshalling to simplify state serialization

Signed-off-by: Thane Thomson <[email protected]>

* e2e: Remove unnecessary method

Signed-off-by: Thane Thomson <[email protected]>

* e2e: Swap Query return value order

Signed-off-by: Thane Thomson <[email protected]>

---------

Signed-off-by: Thane Thomson <[email protected]>

Author: thanethomson

test/e2e/app/app.go

@@ -125,12 +125,13 @@ func NewApplication(cfg *Config) (*Application, error) {
 }
 
 // Info implements ABCI.
-func (app *Application) Info(_ context.Context, _ *abci.RequestInfo) (*abci.ResponseInfo, error) {
+func (app *Application) Info(context.Context, *abci.RequestInfo) (*abci.ResponseInfo, error) {
+	height, hash := app.state.Info()
 	return &abci.ResponseInfo{
 		Version:          version.ABCIVersion,
 		AppVersion:       appVersion,
-		LastBlockHeight:  int64(app.state.Height),
-		LastBlockAppHash: app.state.Hash,
+		LastBlockHeight:  int64(height),
+		LastBlockAppHash: hash,
 	}, nil
 }
 
@@ -160,7 +161,7 @@ func (app *Application) InitChain(_ context.Context, req *abci.RequestInitChain)
 		}
 	}
 	resp := &abci.ResponseInitChain{
-		AppHash: app.state.Hash,
+		AppHash: app.state.GetHash(),
 	}
 	if resp.Validators, err = app.validatorUpdates(0); err != nil {
 		return nil, err
@@ -261,15 +262,16 @@ func (app *Application) Commit(_ context.Context, _ *abci.RequestCommit) (*abci.
 
 // Query implements ABCI.
 func (app *Application) Query(_ context.Context, req *abci.RequestQuery) (*abci.ResponseQuery, error) {
+	value, height := app.state.Query(string(req.Data))
 	return &abci.ResponseQuery{
-		Height: int64(app.state.Height),
+		Height: int64(height),
 		Key:    req.Data,
-		Value:  []byte(app.state.Get(string(req.Data))),
+		Value:  []byte(value),
 	}, nil
 }
 
 // ListSnapshots implements ABCI.
-func (app *Application) ListSnapshots(_ context.Context, _ *abci.RequestListSnapshots) (*abci.ResponseListSnapshots, error) {
+func (app *Application) ListSnapshots(context.Context, *abci.RequestListSnapshots) (*abci.ResponseListSnapshots, error) {
 	snapshots, err := app.snapshots.List()
 	if err != nil {
 		panic(err)
@@ -493,7 +495,7 @@ func (app *Application) Rollback() error {
 }
 
 func (app *Application) getAppHeight() int64 {
-	initialHeightStr := app.state.Get(prefixReservedKey + suffixInitialHeight)
+	initialHeightStr, height := app.state.Query(prefixReservedKey + suffixInitialHeight)
 	if len(initialHeightStr) == 0 {
 		panic("initial height not set in database")
 	}
@@ -502,7 +504,7 @@ func (app *Application) getAppHeight() int64 {
 		panic(fmt.Errorf("malformed initial height %q in database", initialHeightStr))
 	}
 
-	appHeight := int64(app.state.Height)
+	appHeight := int64(height)
 	if appHeight == 0 {
 		appHeight = initialHeight - 1
 	}

test/e2e/app/snapshots.go

@@ -31,7 +31,7 @@ type SnapshotStore struct {
 // NewSnapshotStore creates a new snapshot store.
 func NewSnapshotStore(dir string) (*SnapshotStore, error) {
 	store := &SnapshotStore{dir: dir}
-	if err := os.MkdirAll(dir, 0755); err != nil {
+	if err := os.MkdirAll(dir, 0o755); err != nil {
 		return nil, err
 	}
 	if err := store.loadMetadata(); err != nil {
@@ -84,17 +84,17 @@ func (s *SnapshotStore) saveMetadata() error {
 func (s *SnapshotStore) Create(state *State) (abci.Snapshot, error) {
 	s.Lock()
 	defer s.Unlock()
-	bz, err := state.Export()
+	bz, height, stateHash, err := state.Export()
 	if err != nil {
 		return abci.Snapshot{}, err
 	}
 	snapshot := abci.Snapshot{
-		Height: state.Height,
+		Height: height,
 		Format: 1,
-		Hash:   hashItems(state.Values, state.Height),
+		Hash:   stateHash,
 		Chunks: byteChunks(bz),
 	}
-	err = os.WriteFile(filepath.Join(s.dir, fmt.Sprintf("%v.json", state.Height)), bz, 0o644) //nolint:gosec
+	err = os.WriteFile(filepath.Join(s.dir, fmt.Sprintf("%v.json", height)), bz, 0o644) //nolint:gosec
 	if err != nil {
 		return abci.Snapshot{}, err
 	}

test/e2e/app/state.go

@@ -17,14 +17,21 @@ const (
 	prevStateFileName = "prev_app_state.json"
 )
 
-// State is the application state.
-type State struct {
-	sync.RWMutex
+// Intermediate type used exclusively in serialization/deserialization of
+// State, such that State need not expose any of its internal values publicly.
+type serializedState struct {
 	Height uint64
 	Values map[string]string
 	Hash   []byte
+}
+
+// State is the application state.
+type State struct {
+	sync.RWMutex
+	height uint64
+	values map[string]string
+	hash   []byte
 
-	// private fields aren't marshaled to disk.
 	currentFile string
 	// app saves current and previous state for rollback functionality
 	previousFile    string
@@ -35,12 +42,12 @@ type State struct {
 // NewState creates a new state.
 func NewState(dir string, persistInterval uint64) (*State, error) {
 	state := &State{
-		Values:          make(map[string]string),
+		values:          make(map[string]string),
 		currentFile:     filepath.Join(dir, stateFileName),
 		previousFile:    filepath.Join(dir, prevStateFileName),
 		persistInterval: persistInterval,
 	}
-	state.Hash = hashItems(state.Values, state.Height)
+	state.hash = hashItems(state.values, state.height)
 	err := state.load()
 	switch {
 	case errors.Is(err, os.ErrNotExist):
@@ -66,8 +73,7 @@ func (s *State) load() error {
 			return fmt.Errorf("failed to read state from %q: %w", s.currentFile, err)
 		}
 	}
-	err = json.Unmarshal(bz, s)
-	if err != nil {
+	if err := json.Unmarshal(bz, s); err != nil {
 		return fmt.Errorf("invalid state data in %q: %w", s.currentFile, err)
 	}
 	return nil
@@ -97,11 +103,39 @@ func (s *State) save() error {
 	return os.Rename(newFile, s.currentFile)
 }
 
+// GetHash provides a thread-safe way of accessing a copy of the current state
+// hash.
+func (s *State) GetHash() []byte {
+	s.RLock()
+	defer s.RUnlock()
+	hash := make([]byte, len(s.hash))
+	copy(hash, s.hash)
+	return hash
+}
+
+// Info returns both the height and hash simultaneously, and is used in the
+// ABCI Info call.
+func (s *State) Info() (uint64, []byte) {
+	s.RLock()
+	defer s.RUnlock()
+	height := s.height
+	hash := make([]byte, len(s.hash))
+	copy(hash, s.hash)
+	return height, hash
+}
+
 // Export exports key/value pairs as JSON, used for state sync snapshots.
-func (s *State) Export() ([]byte, error) {
+// Additionally returns the current height and hash of the state.
+func (s *State) Export() ([]byte, uint64, []byte, error) {
 	s.RLock()
 	defer s.RUnlock()
-	return json.Marshal(s.Values)
+	bz, err := json.Marshal(s.values)
+	if err != nil {
+		return nil, 0, nil, err
+	}
+	height := s.height
+	stateHash := hashItems(s.values, height)
+	return bz, height, stateHash, nil
 }
 
 // Import imports key/value pairs from JSON bytes, used for InitChain.AppStateBytes and
@@ -114,71 +148,100 @@ func (s *State) Import(height uint64, jsonBytes []byte) error {
 	if err != nil {
 		return fmt.Errorf("failed to decode imported JSON data: %w", err)
 	}
-	s.Height = height
-	s.Values = values
-	s.Hash = hashItems(values, height)
+	s.height = height
+	s.values = values
+	s.hash = hashItems(values, height)
 	return s.save()
 }
 
 // Get fetches a value. A missing value is returned as an empty string.
 func (s *State) Get(key string) string {
 	s.RLock()
 	defer s.RUnlock()
-	return s.Values[key]
+	return s.values[key]
 }
 
 // Set sets a value. Setting an empty value is equivalent to deleting it.
 func (s *State) Set(key, value string) {
 	s.Lock()
 	defer s.Unlock()
 	if value == "" {
-		delete(s.Values, key)
+		delete(s.values, key)
 	} else {
-		s.Values[key] = value
+		s.values[key] = value
 	}
 }
 
+// Query is used in the ABCI Query call, and provides both the current height
+// and the value associated with the given key.
+func (s *State) Query(key string) (string, uint64) {
+	s.RLock()
+	defer s.RUnlock()
+	height := s.height
+	value := s.values[key]
+	return value, height
+}
+
 // Finalize is called after applying a block, updating the height and returning the new app_hash
 func (s *State) Finalize() []byte {
 	s.Lock()
 	defer s.Unlock()
 	switch {
-	case s.Height > 0:
-		s.Height++
+	case s.height > 0:
+		s.height++
 	case s.initialHeight > 0:
-		s.Height = s.initialHeight
+		s.height = s.initialHeight
 	default:
-		s.Height = 1
+		s.height = 1
 	}
-	s.Hash = hashItems(s.Values, s.Height)
-	return s.Hash
+	s.hash = hashItems(s.values, s.height)
+	return s.hash
 }
 
 // Commit commits the current state.
 func (s *State) Commit() (uint64, error) {
 	s.Lock()
 	defer s.Unlock()
-	if s.persistInterval > 0 && s.Height%s.persistInterval == 0 {
+	if s.persistInterval > 0 && s.height%s.persistInterval == 0 {
 		err := s.save()
 		if err != nil {
 			return 0, err
 		}
 	}
-	return s.Height, nil
+	return s.height, nil
 }
 
 func (s *State) Rollback() error {
 	bz, err := os.ReadFile(s.previousFile)
 	if err != nil {
 		return fmt.Errorf("failed to read state from %q: %w", s.previousFile, err)
 	}
-	err = json.Unmarshal(bz, s)
-	if err != nil {
+	if err := json.Unmarshal(bz, s); err != nil {
 		return fmt.Errorf("invalid state data in %q: %w", s.previousFile, err)
 	}
 	return nil
 }
 
+func (s *State) UnmarshalJSON(b []byte) error {
+	var ss serializedState
+	if err := json.Unmarshal(b, &ss); err != nil {
+		return err
+	}
+	s.height = ss.Height
+	s.values = ss.Values
+	s.hash = ss.Hash
+	return nil
+}
+
+func (s *State) MarshalJSON() ([]byte, error) {
+	ss := &serializedState{
+		Height: s.height,
+		Values: s.values,
+		Hash:   s.hash,
+	}
+	return json.Marshal(ss)
+}
+
 // hashItems hashes a set of key/value items.
 func hashItems(items map[string]string, height uint64) []byte {
 	keys := make([]string, 0, len(items))