GH-1328: Improve behavior of Dictionary iterator invalidation

Previously, an assertion was triggered in debug builds upon any attempt
to insert or remove a Dictionary entry while any iteration of that
Dictionary is underway and also even in cases where Dictionary membership
was not actually modified (and thus invalidates a loop).

Now, it emits run-time warnings regardless of build-type and only when
insert/remove operations truly change the Dictionary membership.  In the
context of a Zeek script causing an invalidation, the warning message
also now helps pinpoint the exact expression that causes it.

src/Dict.cc

@@ -217,6 +217,77 @@ TEST_CASE("dict iteration")
 	delete key2;
 	}
 
+TEST_CASE("dict iterator invalidation")
+	{
+	PDict<uint32_t> dict;
+
+	uint32_t val = 15;
+	uint32_t key_val = 5;
+	auto key = new detail::HashKey(key_val);
+
+	uint32_t val2 = 10;
+	uint32_t key_val2 = 25;
+	auto key2 = new detail::HashKey(key_val2);
+
+	uint32_t val3 = 42;
+	uint32_t key_val3 = 37;
+	auto key3 = new detail::HashKey(key_val3);
+
+	dict.Insert(key, &val);
+	dict.Insert(key2, &val2);
+
+	detail::HashKey* it_key;
+	bool iterators_invalidated = false;
+	IterCookie* it = dict.InitForIteration();
+	CHECK(it != nullptr);
+
+	while ( uint32_t* entry = dict.NextEntry(it_key, it) )
+		{
+		iterators_invalidated = false;
+		dict.Remove(key3, &iterators_invalidated);
+		// Key doesn't exist, nothing to remove, iteration not invalidated.
+		CHECK(!iterators_invalidated);
+
+		iterators_invalidated = false;
+		dict.Insert(key, &val2, &iterators_invalidated);
+		// Key exists, value gets overwritten, iteration not invalidated.
+		CHECK(!iterators_invalidated);
+
+		iterators_invalidated = false;
+		dict.Remove(key2, &iterators_invalidated);
+		// Key exists, gets removed, iteration is invalidated.
+		CHECK(iterators_invalidated);
+
+		delete it_key;
+		dict.StopIteration(it);
+		break;
+		}
+
+	it = dict.InitForIteration();
+	CHECK(it != nullptr);
+
+	while ( uint32_t* entry = dict.NextEntry(it_key, it) )
+		{
+		iterators_invalidated = false;
+		dict.Insert(key3, &val3, &iterators_invalidated);
+		// Key doesn't exist, gets inserted, iteration is invalidated.
+		CHECK(iterators_invalidated);
+
+		delete it_key;
+		dict.StopIteration(it);
+		break;
+		}
+
+	CHECK(dict.Length() == 2);
+	CHECK(*static_cast<uint32_t*>(dict.Lookup(key)) == val2);
+	CHECK(*static_cast<uint32_t*>(dict.Lookup(key3)) == val3);
+	CHECK(static_cast<uint32_t*>(dict.Lookup(key2)) == nullptr);
+
+	delete key;
+	delete key2;
+	delete key3;
+	}
+
 TEST_SUITE_END();
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
@@ -718,13 +789,10 @@ int Dictionary::LookupIndex(const void* key, int key_size, detail::hash_t hash,
 // Insert
 //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
 
-void* Dictionary::Insert(void* key, int key_size, detail::hash_t hash, void* val, bool copy_key)
+void* Dictionary::Insert(void* key, int key_size, detail::hash_t hash, void* val, bool copy_key, bool* iterators_invalidated)
 	{
 	ASSERT_VALID(this);
 
-	// Allow insertions only if there's no active non-robust iterations.
-	ASSERT(num_iterators == 0 || (cookies && cookies->size() == num_iterators));
-
 	// Initialize the table if it hasn't been done yet. This saves memory storing a bunch
 	// of empty dicts.
 	if ( ! table )
@@ -762,6 +830,14 @@ void* Dictionary::Insert(void* key, int key_size, detail::hash_t hash, void* val
 		}
 	else
 		{
+		if ( ! HaveOnlyRobustIterators() )
+			{
+			if ( iterators_invalidated )
+				*iterators_invalidated = true;
+			else
+				reporter->InternalWarning("Dictionary::Insert() possibly caused iterator invalidation");
+			}
+
 		// Allocate memory for key if necesary. Key is updated to reflect internal key if necessary.
 		detail::DictEntry entry(key, key_size, hash, val, insert_distance, copy_key);
 		InsertRelocateAndAdjust(entry, insert_position);
@@ -879,16 +955,24 @@ void Dictionary::SizeUp()
 // Remove
 /////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
 
-void* Dictionary::Remove(const void* key, int key_size, detail::hash_t hash, bool dont_delete)
+void* Dictionary::Remove(const void* key, int key_size, detail::hash_t hash, bool dont_delete, bool* iterators_invalidated)
 	{//cookie adjustment: maintain inserts here. maintain next in lower level version.
 	ASSERT_VALID(this);
-	ASSERT(num_iterators == 0 || (cookies && cookies->size() == num_iterators)); //only robust iterators exist.
+
 	ASSERT(! dont_delete); //this is a poorly designed flag. if on, the internal has nowhere to return and memory is lost.
 
 	int position = LookupIndex(key, key_size, hash);
 	if ( position < 0 )
 		return nullptr;
 
+	if ( ! HaveOnlyRobustIterators() )
+		{
+		if ( iterators_invalidated )
+			*iterators_invalidated = true;
+		else
+			reporter->InternalWarning("Dictionary::Remove() possibly caused iterator invalidation");
+		}
+
 	detail::DictEntry entry = RemoveRelocateAndAdjust(position);
 	num_entries--;
 	ASSERT(num_entries >= 0);

src/Dict.h

@@ -164,20 +164,26 @@ public:
 	void* Lookup(const void* key, int key_size, detail::hash_t h) const;
 
 	// Returns previous value, or 0 if none.
-	void* Insert(detail::HashKey* key, void* val)
-		{ return Insert(key->TakeKey(), key->Size(), key->Hash(), val, false); }
+	// If iterators_invalidated is supplied, its value is set to true
+	// if the removal may have invalidated any existing iterators.
+	void* Insert(detail::HashKey* key, void* val, bool* iterators_invalidated = nullptr)
+		{ return Insert(key->TakeKey(), key->Size(), key->Hash(), val, false, iterators_invalidated); }
 
 	// If copy_key is true, then the key is copied, otherwise it's assumed
 	// that it's a heap pointer that now belongs to the Dictionary to
 	// manage as needed.
-	void* Insert(void* key, int key_size, detail::hash_t hash, void* val, bool copy_key);
+	// If iterators_invalidated is supplied, its value is set to true
+	// if the removal may have invalidated any existing iterators.
+	void* Insert(void* key, int key_size, detail::hash_t hash, void* val, bool copy_key, bool* iterators_invalidated = nullptr);
 
 	// Removes the given element.  Returns a pointer to the element in
 	// case it needs to be deleted.  Returns 0 if no such element exists.
 	// If dontdelete is true, the key's bytes will not be deleted.
-	void* Remove(const detail::HashKey* key)
-		{ return Remove(key->Key(), key->Size(), key->Hash()); }
-	void* Remove(const void* key, int key_size, detail::hash_t hash, bool dont_delete = false);
+	// If iterators_invalidated is supplied, its value is set to true
+	// if the removal may have invalidated any existing iterators.
+	void* Remove(const detail::HashKey* key, bool* iterators_invalidated = nullptr)
+		{ return Remove(key->Key(), key->Size(), key->Hash(), false, iterators_invalidated); }
+	void* Remove(const void* key, int key_size, detail::hash_t hash, bool dont_delete = false, bool* iterators_invalidated = nullptr);
 
 	// Number of entries.
 	int Length() const
@@ -337,6 +343,9 @@ private:
 
 	void SizeUp();
 
+	bool HaveOnlyRobustIterators() const
+		{ return num_iterators == 0 || (cookies && cookies->size() == num_iterators); }
+
 	//alligned on 8-bytes with 4-leading bytes. 7*8=56 bytes a dictionary.
 
 	// when sizeup but the current mapping is in progress. the current mapping will be ignored
@@ -380,13 +389,13 @@ public:
 		}
 	T* Lookup(const detail::HashKey* key) const
 		{ return (T*) Dictionary::Lookup(key); }
-	T* Insert(const char* key, T* val)
+	T* Insert(const char* key, T* val, bool* iterators_invalidated = nullptr)
 		{
 		detail::HashKey h(key);
-		return (T*) Dictionary::Insert(&h, (void*) val);
+		return (T*) Dictionary::Insert(&h, (void*) val, iterators_invalidated);
 		}
-	T* Insert(detail::HashKey* key, T* val)
-		{ return (T*) Dictionary::Insert(key, (void*) val); }
+	T* Insert(detail::HashKey* key, T* val, bool* iterators_invalidated = nullptr)
+		{ return (T*) Dictionary::Insert(key, (void*) val, iterators_invalidated); }
 	T* NthEntry(int n) const
 		{ return (T*) Dictionary::NthEntry(n); }
 	T* NthEntry(int n, const char*& key) const
@@ -401,10 +410,10 @@ public:
 		}
 	T* NextEntry(detail::HashKey*& h, IterCookie*& cookie) const
 		{ return (T*) Dictionary::NextEntry(h, cookie, true); }
-	T* RemoveEntry(const detail::HashKey* key)
-		{ return (T*) Remove(key->Key(), key->Size(), key->Hash()); }
-	T* RemoveEntry(const detail::HashKey& key)
-		{ return (T*) Remove(key.Key(), key.Size(), key.Hash()); }
+	T* RemoveEntry(const detail::HashKey* key, bool* iterators_invalidated = nullptr)
+		{ return (T*) Remove(key->Key(), key->Size(), key->Hash(), false, iterators_invalidated); }
+	T* RemoveEntry(const detail::HashKey& key, bool* iterators_invalidated = nullptr)
+		{ return (T*) Remove(key.Key(), key.Size(), key.Hash(), false, iterators_invalidated); }
 };
 
 } // namespace zeek

src/Expr.cc

@@ -2619,7 +2619,17 @@ void IndexExpr::Add(Frame* f)
 	if ( ! v2 )
 		return;
 
-	v1->AsTableVal()->Assign(std::move(v2), nullptr);
+	bool iterators_invalidated = false;
+	v1->AsTableVal()->Assign(std::move(v2), nullptr, true, &iterators_invalidated);
+
+	if ( iterators_invalidated )
+		{
+		ODesc d;
+		Describe(&d);
+		reporter->PushLocation(GetLocationInfo());
+		reporter->Warning("possible loop/iterator invalidation caused by expression: %s", d.Description());
+		reporter->PopLocation();
+		}
 	}
 
 void IndexExpr::Delete(Frame* f)
@@ -2637,7 +2647,17 @@ void IndexExpr::Delete(Frame* f)
 	if ( ! v2 )
 		return;
 
-	v1->AsTableVal()->Remove(*v2);
+	bool iterators_invalidated = false;
+	v1->AsTableVal()->Remove(*v2, true, &iterators_invalidated);
+
+	if ( iterators_invalidated )
+		{
+		ODesc d;
+		Describe(&d);
+		reporter->PushLocation(GetLocationInfo());
+		reporter->Warning("possible loop/iterator invalidation caused by expression: %s", d.Description());
+		reporter->PopLocation();
+		}
 	}
 
 ExprPtr IndexExpr::MakeLvalue()
@@ -2858,7 +2878,10 @@ void IndexExpr::Assign(Frame* f, ValPtr v)
 		}
 
 	case TYPE_TABLE:
-		if ( ! v1->AsTableVal()->Assign(std::move(v2), std::move(v)) )
+		{
+		bool iterators_invalidated = false;
+
+		if ( ! v1->AsTableVal()->Assign(std::move(v2), std::move(v), true, &iterators_invalidated) )
 			{
 			v = std::move(v_extra);
 
@@ -2876,6 +2899,16 @@ void IndexExpr::Assign(Frame* f, ValPtr v)
 			else
 				RuntimeErrorWithCallStack("assignment failed with null value");
 			}
+
+		if ( iterators_invalidated )
+			{
+			ODesc d;
+			Describe(&d);
+			reporter->PushLocation(GetLocationInfo());
+			reporter->Warning("possible loop/iterator invalidation caused by expression: %s", d.Description());
+			reporter->PopLocation();
+			}
+		}
 		break;
 
 	case TYPE_STRING:

src/Val.cc

@@ -1548,7 +1548,8 @@ void TableVal::CheckExpireAttr(detail::AttrTag at)
 		}
 	}
 
-bool TableVal::Assign(ValPtr index, ValPtr new_val, bool broker_forward)
+bool TableVal::Assign(ValPtr index, ValPtr new_val, bool broker_forward,
+                      bool* iterators_invalidated)
 	{
 	auto k = MakeHashKey(*index);
 
@@ -1558,7 +1559,7 @@ bool TableVal::Assign(ValPtr index, ValPtr new_val, bool broker_forward)
 		return false;
 		}
 
-	return Assign(std::move(index), std::move(k), std::move(new_val), broker_forward);
+	return Assign(std::move(index), std::move(k), std::move(new_val), broker_forward, iterators_invalidated);
 	}
 
 bool TableVal::Assign(Val* index, Val* new_val)
@@ -1567,7 +1568,7 @@ bool TableVal::Assign(Val* index, Val* new_val)
 	}
 
 bool TableVal::Assign(ValPtr index, std::unique_ptr<detail::HashKey> k,
-                      ValPtr new_val, bool broker_forward)
+                      ValPtr new_val, bool broker_forward, bool* iterators_invalidated)
 	{
 	bool is_set = table_type->IsSet();
 
@@ -1576,7 +1577,7 @@ bool TableVal::Assign(ValPtr index, std::unique_ptr<detail::HashKey> k,
 
 	TableEntryVal* new_entry_val = new TableEntryVal(std::move(new_val));
 	detail::HashKey k_copy(k->Key(), k->Size(), k->Hash());
-	TableEntryVal* old_entry_val = AsNonConstTable()->Insert(k.get(), new_entry_val);
+	TableEntryVal* old_entry_val = AsNonConstTable()->Insert(k.get(), new_entry_val, iterators_invalidated);
 
 	// If the dictionary index already existed, the insert may free up the
 	// memory allocated to the key bytes, so have to assume k is invalid
@@ -2263,11 +2264,11 @@ void TableVal::SendToStore(const Val* index, const TableEntryVal* new_entry_val,
 		}
 	}
 
-ValPtr TableVal::Remove(const Val& index, bool broker_forward)
+ValPtr TableVal::Remove(const Val& index, bool broker_forward, bool* iterators_invalidated)
 	{
 	auto k = MakeHashKey(index);
 
-	TableEntryVal* v = k ? AsNonConstTable()->RemoveEntry(k.get()) : nullptr;
+	TableEntryVal* v = k ? AsNonConstTable()->RemoveEntry(k.get(), iterators_invalidated) : nullptr;
 	ValPtr va;
 
 	if ( v )
@@ -2293,9 +2294,9 @@ ValPtr TableVal::Remove(const Val& index, bool broker_forward)
 	return va;
 	}
 
-ValPtr TableVal::Remove(const detail::HashKey& k)
+ValPtr TableVal::Remove(const detail::HashKey& k, bool* iterators_invalidated)
 	{
-	TableEntryVal* v = AsNonConstTable()->RemoveEntry(k);
+	TableEntryVal* v = AsNonConstTable()->RemoveEntry(k, iterators_invalidated);
 	ValPtr va;
 
 	if ( v )

src/Val.h

@@ -792,9 +792,12 @@ public:
 	 * must be nullptr.
 	 * @param broker_forward Controls if the value will be forwarded to attached
 	 *        Broker stores.
+	 * @param iterators_invalidated  if supplied, gets set to true if the operation
+	 *        may have invalidated existing iterators.
 	 * @return  True if the assignment type-checked.
 	 */
-	bool Assign(ValPtr index, ValPtr new_val, bool broker_forward = true);
+	bool Assign(ValPtr index, ValPtr new_val, bool broker_forward = true,
+	            bool* iterators_invalidated = nullptr);
 
 	/**
 	 * Assigns a value at an associated index in the table (or in the
@@ -803,13 +806,16 @@ public:
 	 * (if needed, the index val can be recovered from the hash key).
 	 * @param k  A precomputed hash key to use.
 	 * @param new_val  The value to assign at the index.  For a set, this
+	 * @param iterators_invalidated  if supplied, gets set to true if the operation
+	 *        may have invalidated existing iterators.
 	 * must be nullptr.
 	 * @param broker_forward Controls if the value will be forwarded to attached
 	 *        Broker stores.
 	 * @return  True if the assignment type-checked.
 	 */
 	bool Assign(ValPtr index, std::unique_ptr<detail::HashKey> k,
-	            ValPtr new_val, bool broker_forward = true);
+	            ValPtr new_val, bool broker_forward = true,
+	            bool* iterators_invalidated = nullptr);
 
 	// Returns true if the assignment typechecked, false if not. The
 	// methods take ownership of new_val, but not of the index.  If we're
@@ -943,19 +949,23 @@ public:
 	 * @param index  The index to remove.
 	 * @param broker_forward Controls if the remove operation will be forwarded to attached
 	 *        Broker stores.
+	 * @param iterators_invalidated  if supplied, gets set to true if the operation
+	 *        may have invalidated existing iterators.
 	 * @return  The value associated with the index if it exists, else nullptr.
 	 * For a sets that don't really contain associated values, a placeholder
 	 * value is returned to differentiate it from non-existent index (nullptr),
 	 * but otherwise has no meaning in relation to the set's contents.
 	 */
-	ValPtr Remove(const Val& index, bool broker_forward = true);
+	ValPtr Remove(const Val& index, bool broker_forward = true, bool* iterators_invalidated = nullptr);
 
 	/**
 	 * Same as Remove(const Val&), but uses a precomputed hash key.
 	 * @param k  The hash key to lookup.
+	 * @param iterators_invalidated  if supplied, gets set to true if the operation
+	 *        may have invalidated existing iterators.
 	 * @return  Same as Remove(const Val&).
 	 */
-	ValPtr Remove(const detail::HashKey& k);
+	ValPtr Remove(const detail::HashKey& k, bool* iterators_invalidated = nullptr);
 
 	[[deprecated("Remove in v4.1.  Use Remove().")]]
 	Val* Delete(const Val* index)

testing/btest/Baseline/language.table-set-iterator-invalidation/out

@@ -0,0 +1,42 @@
+### BTest baseline data generated by btest-diff. Do not edit. Use "btest -U/-u" to update. Requires BTest >= 0.63.
+warning in <...>/table-set-iterator-invalidation.zeek, line 21: possible loop/iterator invalidation caused by expression: t[4]
+warning in <...>/table-set-iterator-invalidation.zeek, line 30: possible loop/iterator invalidation caused by expression: t[4]
+warning in <...>/table-set-iterator-invalidation.zeek, line 53: possible loop/iterator invalidation caused by expression: s[4]
+warning in <...>/table-set-iterator-invalidation.zeek, line 62: possible loop/iterator invalidation caused by expression: s[4]
+{
+[2] = 2,
+[1] = 1,
+[3] = 3
+}
+{
+[2] = 2,
+[4] = four,
+[3] = 3,
+[1] = 1
+}
+{
+[2] = 2,
+[1] = 1,
+[3] = 3
+}
+{
+[2] = 2,
+[1] = 1,
+[3] = 3
+}
+{
+2,
+4,
+3,
+1
+}
+{
+2,
+1,
+3
+}
+{
+2,
+1,
+3
+}

testing/btest/language/table-set-iterator-invalidation.zeek

@@ -0,0 +1,73 @@
+# @TEST-EXEC: zeek -b %INPUT >out 2>&1
+# @TEST-EXEC: TEST_DIFF_CANONIFIER=$SCRIPTS/diff-remove-abspath btest-diff out
+
+# Note that while modifying container membership during for-loop iteration is
+# supposed to be undefined-behavior, it should be practically ok to have this
+# test perform such operations if they always `break` out of the loop
+# immediately afterward.
+
+local t = table([1] = "one", [2] = "two", [3] = "three");
+
+for ( i in t )
+	# Modifying an existing element is not qualified an modifying membership,
+	# so this doesn't trigger a warning.
+	t[i] = cat(i);
+
+print t;
+
+for ( i in t )
+	{
+	# Adding an element in a loop should trigger a warning.
+	t[4] = "four";
+	break;
+	}
+
+print t;
+
+for ( i in t )
+	{
+	# Deleting an element in a loop should trigger a warning.
+	delete t[4];
+	break;
+	}
+
+print t;
+
+for ( i in t )
+	# Trying to delete a non-existent element within in a loop does not
+	# actually modify membership, so does not trigger a warning.
+	delete t[0];
+
+print t;
+
+local s = set(1, 2, 3);
+
+for ( n in s )
+	# Trying to add an existing element within in a loop does not
+	# actually modify membership, so does not trigger a warning.
+	add s[1];
+
+for ( n in s )
+	{
+	# Adding an element in a loop should trigger a warning.
+	add s[4];
+	break;
+	}
+
+print s;
+
+for ( n in s )
+	{
+	# Deleting an element in a loop should trigger a warning.
+	delete s[4];
+	break;
+	}
+
+print s;
+
+for ( n in s )
+	# Trying to delete a non-existent element within in a loop does not
+	# actually modify membership, so does not trigger a warning.
+	delete s[0];
+
+print s;