Extract methods for accessing EtcdState fields

Signed-off-by: Marek Siarkowicz <siarkowicz@google.com>

Author: serathius

tests/robustness/model/describe.go

@@ -16,9 +16,6 @@ package model
 
 import (
 	"fmt"
-	"maps"
-	"slices"
-	"sort"
 	"strings"
 
 	"go.etcd.io/etcd/client/pkg/v3/types"
@@ -91,19 +88,18 @@ func describeEtcdState(state EtcdState) string {
 	if len(state.KeyValues) > 0 {
 		descHTML = append(descHTML, "keys: <ul style=\"margin: 0.25em 0;\">")
 
-		keys := slices.Collect(maps.Keys(state.KeyValues))
-		sort.Strings(keys)
-		for _, key := range keys {
-			descHTML = append(descHTML, fmt.Sprintf("<li style=\"margin: 0.25em 0;\"><strong>%s</strong> - ", key))
+		keys, values, leases := state.KeysValueLeases()
+		for i := range keys {
+			descHTML = append(descHTML, fmt.Sprintf("<li style=\"margin: 0.25em 0;\"><strong>%s</strong> - ", keys[i]))
 
-			value := state.KeyValues[key]
+			value := values[i]
 			if value.Value.Value != "" {
 				descHTML = append(descHTML, fmt.Sprintf("val: %q, ", value.Value.Value))
 			}
 			if value.Value.Hash != 0 {
 				descHTML = append(descHTML, fmt.Sprintf("hash: %d, ", value.Value.Hash))
 			}
-			lease := state.KeyLeases[key]
+			lease := leases[i]
 			if lease != 0 {
 				descHTML = append(descHTML, fmt.Sprintf("lease: %d, ", lease))
 			}
@@ -114,10 +110,9 @@ func describeEtcdState(state EtcdState) string {
 		descHTML = append(descHTML, "</ul>")
 	}
 
-	if len(state.Leases) > 0 {
+	leases := state.leases()
+	if len(leases) > 0 {
 		descHTML = append(descHTML, "leases: <ul style=\"margin: 0.25em 0;\">")
-		leases := slices.Collect(maps.Keys(state.Leases))
-		slices.Sort(leases)
 		for _, lease := range leases {
 			descHTML = append(descHTML, fmt.Sprintf("<li style=\"margin: 0.25em 0;\"><strong>%d</strong></li>", lease))
 		}

tests/robustness/model/deterministic.go

@@ -20,6 +20,7 @@ import (
 	"html"
 	"maps"
 	"reflect"
+	"slices"
 	"sort"
 
 	"github.com/anishathalye/porcupine"
@@ -163,7 +164,10 @@ func (s EtcdState) stepTxn(request EtcdRequest) (EtcdState, MaybeEtcdResponse) {
 	newState := s.DeepCopy()
 	failure := false
 	for _, cond := range request.Txn.Conditions {
-		val := newState.KeyValues[cond.Key]
+		val, ok := newState.GetValue(cond.Key)
+		if !ok {
+			val = &ValueRevision{}
+		}
 		if cond.ExpectedVersion > 0 {
 			if val.Version != cond.ExpectedVersion {
 				failure = true
@@ -187,29 +191,29 @@ func (s EtcdState) stepTxn(request EtcdRequest) (EtcdState, MaybeEtcdResponse) {
 				RangeResponse: newState.getRange(op.Range),
 			}
 		case PutOperation:
-			_, leaseExists := newState.Leases[op.Put.LeaseID]
-			if op.Put.LeaseID != 0 && !leaseExists {
-				break
+			var leaseID *int64
+			if op.Put.LeaseID != 0 {
+				if !newState.leaseExists(op.Put.LeaseID) {
+					break
+				}
+				leaseID = &op.Put.LeaseID
 			}
 			ver := int64(1)
-			if val, exists := newState.KeyValues[op.Put.Key]; exists && val.Version > 0 {
-				ver = val.Version + 1
+			valPtr, exists := newState.GetValue(op.Put.Key)
+			if exists && valPtr.Version > 0 {
+				ver = valPtr.Version + 1
 			}
-			newState.KeyValues[op.Put.Key] = ValueRevision{
+			val := ValueRevision{
 				Value:       op.Put.Value,
 				ModRevision: newState.Revision + 1,
 				Version:     ver,
 			}
+			newState.setValueLease(op.Put.Key, val, leaseID)
 			increaseRevision = true
-			newState = detachFromOldLease(newState, op.Put.Key)
-			if leaseExists {
-				newState = attachToNewLease(newState, op.Put.LeaseID, op.Put.Key)
-			}
 		case DeleteOperation:
-			if _, ok := newState.KeyValues[op.Delete.Key]; ok {
-				delete(newState.KeyValues, op.Delete.Key)
+			if _, ok := newState.GetValue(op.Delete.Key); ok {
+				newState.deleteKey(op.Delete.Key)
 				increaseRevision = true
-				newState = detachFromOldLease(newState, op.Delete.Key)
 				opResp[i].Deleted = 1
 			}
 		default:
@@ -294,28 +298,69 @@ func (s EtcdState) getRange(options RangeOptions) RangeResponse {
 		}
 		response.Count = count
 	} else {
-		value, ok := s.KeyValues[options.Start]
+		valPtr, ok := s.GetValue(options.Start)
 		if ok {
 			response.KVs = append(response.KVs, KeyValue{
 				Key:           options.Start,
-				ValueRevision: value,
+				ValueRevision: *valPtr,
 			})
 			response.Count = 1
 		}
 	}
 	return response
 }
 
-func detachFromOldLease(s EtcdState, key string) EtcdState {
+func (s EtcdState) KeysValueLeases() (keys []string, values []ValueRevision, leases []int64) {
+	keys = make([]string, 0, len(s.KeyValues))
+	values = make([]ValueRevision, 0, len(s.KeyValues))
+	leases = make([]int64, 0, len(s.KeyLeases))
+
+	for k, v := range s.KeyValues {
+		keys = append(keys, k)
+		values = append(values, v)
+		leases = append(leases, s.KeyLeases[k])
+	}
+	return keys, values, leases
+}
+
+func (s EtcdState) leases() []int64 {
+	return slices.Collect(maps.Keys(s.Leases))
+}
+
+func (s EtcdState) GetValue(key string) (*ValueRevision, bool) {
+	val, ok := s.KeyValues[key]
+	if !ok {
+		return nil, false
+	}
+	return &val, true
+}
+
+func (s EtcdState) setValueLease(key string, val ValueRevision, lease *int64) {
+	s.KeyValues[key] = val
 	if oldLeaseID, ok := s.KeyLeases[key]; ok {
 		delete(s.Leases[oldLeaseID].Keys, key)
+	}
+	if lease != nil {
+		s.KeyLeases[key] = *lease
+		s.Leases[*lease].Keys[key] = leased
+	} else {
 		delete(s.KeyLeases, key)
 	}
-	return s
 }
 
-func attachToNewLease(s EtcdState, leaseID int64, key string) EtcdState {
-	s.KeyLeases[key] = leaseID
-	s.Leases[leaseID].Keys[key] = leased
-	return s
+func (s EtcdState) leaseExists(lease int64) bool {
+	_, ok := s.Leases[lease]
+	return ok
+}
+
+func (s EtcdState) deleteKey(key string) {
+	delete(s.KeyValues, key)
+	if oldLeaseID, ok := s.KeyLeases[key]; ok {
+		delete(s.Leases[oldLeaseID].Keys, key)
+	}
+	delete(s.KeyLeases, key)
+}
+
+func (s EtcdState) leaseKeys(leaseID int64) []string {
+	return slices.Sorted(maps.Keys(s.Leases[leaseID].Keys))
 }

tests/robustness/model/replay.go

@@ -16,7 +16,6 @@ package model
 
 import (
 	"fmt"
-	"sort"
 	"strings"
 
 	"github.com/anishathalye/porcupine"
@@ -99,7 +98,7 @@ func toWatchEvents(prevState *EtcdState, request EtcdRequest, response MaybeEtcd
 					events = append(events, e)
 				}
 			case PutOperation:
-				_, leaseExists := prevState.Leases[op.Put.LeaseID]
+				leaseExists := prevState.leaseExists(op.Put.LeaseID)
 				if op.Put.LeaseID != 0 && !leaseExists {
 					break
 				}
@@ -112,7 +111,7 @@ func toWatchEvents(prevState *EtcdState, request EtcdRequest, response MaybeEtcd
 					},
 					Revision: response.Revision,
 				}
-				if _, ok := prevState.KeyValues[op.Put.Key]; !ok {
+				if _, ok := prevState.GetValue(op.Put.Key); !ok {
 					e.IsCreate = true
 				}
 				events = append(events, e)
@@ -121,15 +120,7 @@ func toWatchEvents(prevState *EtcdState, request EtcdRequest, response MaybeEtcd
 			}
 		}
 	case LeaseRevoke:
-		deletedKeys := []string{}
-		for key := range prevState.Leases[request.LeaseRevoke.LeaseID].Keys {
-			if _, ok := prevState.KeyValues[key]; ok {
-				deletedKeys = append(deletedKeys, key)
-			}
-		}
-
-		sort.Strings(deletedKeys)
-		for _, key := range deletedKeys {
+		for _, key := range prevState.leaseKeys(request.LeaseRevoke.LeaseID) {
 			e := PersistedEvent{
 				Event: Event{
 					Type: DeleteOperation,

tests/robustness/validate/validate_test.go

@@ -226,6 +226,39 @@ func TestValidateWatch(t *testing.T) {
 				putRequest("b", "2"),
 			},
 		},
+		{
+			name: "Reliable - Put with non-existent lease doesn't generate watch event - pass",
+			reports: []report.ClientReport{
+				{
+					Watch: []model.WatchOperation{
+						{
+							Request: model.WatchRequest{
+								Key: "a",
+							},
+							Responses: []model.WatchResponse{},
+						},
+					},
+				},
+			},
+			persistedRequests: []model.EtcdRequest{
+				{
+					Type: model.Txn,
+					Txn: &model.TxnRequest{
+						OperationsOnSuccess: []model.EtcdOperation{
+							{
+								Type: model.PutOperation,
+								Put: model.PutOptions{
+									Key:     "a",
+									Value:   model.ToValueOrHash("1"),
+									LeaseID: 1,
+								},
+							},
+						},
+					},
+				},
+			},
+			expectError: "",
+		},
 		{
 			name: "Ordered, Unique - unique ordered events in separate response - pass",
 			reports: []report.ClientReport{

tests/robustness/validate/watch.go

@@ -278,9 +278,15 @@ func validatePrevKV(lg *zap.Logger, replay *model.EtcdReplay, report report.Clie
 
 				// We allow PrevValue to be nil since in the face of compaction, etcd does not
 				// guarantee its presence.
-				if event.PrevValue != nil && *event.PrevValue != state.KeyValues[event.Key] {
-					lg.Error("Incorrect event prevValue field", zap.Int("client", report.ClientID), zap.Any("event", event), zap.Any("previousValue", state.KeyValues[event.Key]))
-					err = errBrokePrevKV
+				if event.PrevValue != nil {
+					val, ok := state.GetValue(event.Key)
+					if !ok {
+						lg.Error("Incorrect event prevValue field", zap.Int("client", report.ClientID), zap.Any("event", event), zap.Any("previousValue", nil))
+						err = errBrokePrevKV
+					} else if *event.PrevValue != *val {
+						lg.Error("Incorrect event prevValue field", zap.Int("client", report.ClientID), zap.Any("event", event), zap.Any("previousValue", *val))
+						err = errBrokePrevKV
+					}
 				}
 			}
 		}
@@ -299,7 +305,7 @@ func validateIsCreate(lg *zap.Logger, replay *model.EtcdReplay, report report.Cl
 				}
 				// A create event will not have an entry in our history and a non-create
 				// event *should* have an entry in our history.
-				if _, prevKeyExists := state.KeyValues[event.Key]; event.IsCreate == prevKeyExists {
+				if _, prevKeyExists := state.GetValue(event.Key); event.IsCreate == prevKeyExists {
 					lg.Error("Incorrect event IsCreate field", zap.Int("client", report.ClientID), zap.Any("event", event))
 					err = errBrokeIsCreate
 				}