From b273ad83ff0bb53fa68012ca38085ecb0eff9749 Mon Sep 17 00:00:00 2001
From: sfc-gh-tclinkenbeard <trevor.clinkenbeard@snowflake.com>
Date: Thu, 15 Sep 2022 16:12:51 -0700
Subject: [PATCH] Avoid tracing unexpected code probes

---
 flow/CodeProbe.cpp            | 14 +++++++-------
 flow/include/flow/CodeProbe.h | 21 ++++++++++++++++++---
 2 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/flow/CodeProbe.cpp b/flow/CodeProbe.cpp
index 425aa8929d..38961acbd6 100644
--- a/flow/CodeProbe.cpp
+++ b/flow/CodeProbe.cpp
@@ -228,12 +228,12 @@ void CodeProbes::traceMissedProbes(Optional<ExecutionContext> context) const {
 		std::tie(iter, std::ignore) = locations.emplace(probe.first, false);
 		iter->second = iter->second || probe.second->wasHit();
 	}
-	for (auto probe : codeProbes) {
-		auto iter = locations.find(probe.first);
+	for (const auto& [loc, probe] : codeProbes) {
+		auto iter = locations.find(loc);
 		ASSERT(iter != locations.end());
-		if (!iter->second) {
+		if (!iter->second && probe->shouldTrace()) {
 			iter->second = true;
-			probe.second->trace(false);
+			probe->trace(false);
 		}
 	}
 }
@@ -280,15 +280,15 @@ void ICodeProbe::printProbesJSON(std::vector<std::string> const& ctxs) {
 // annotations
 namespace assert {
 
-bool NoSim::operator()(ICodeProbe* self) const {
+bool NoSim::operator()(ICodeProbe const* self) const {
 	return !g_network->isSimulated();
 }
 
-bool SimOnly::operator()(ICodeProbe* self) const {
+bool SimOnly::operator()(ICodeProbe const* self) const {
 	return g_network->isSimulated();
 }
 
-bool RocksDB::operator()(ICodeProbe* self) const {
+bool RocksDB::operator()(ICodeProbe const* self) const {
 #ifdef SSD_ROCKSDB_EXPERIMENTAL
 	return true;
 #else
diff --git a/flow/include/flow/CodeProbe.h b/flow/include/flow/CodeProbe.h
index cf0ba436f5..e4c4e3ea38 100644
--- a/flow/include/flow/CodeProbe.h
+++ b/flow/include/flow/CodeProbe.h
@@ -64,15 +64,15 @@ operator|(Left const& lhs, Right const& rhs) {
 namespace assert {
 struct NoSim {
 	constexpr static AnnotationType type = AnnotationType::Assertion;
-	bool operator()(ICodeProbe* self) const;
+	bool operator()(ICodeProbe const* self) const;
 };
 struct SimOnly {
 	constexpr static AnnotationType type = AnnotationType::Assertion;
-	bool operator()(ICodeProbe* self) const;
+	bool operator()(ICodeProbe const* self) const;
 };
 struct RocksDB {
 	constexpr static AnnotationType type = AnnotationType::Assertion;
-	bool operator()(ICodeProbe* self) const;
+	bool operator()(ICodeProbe const* self) const;
 };
 
 template <class Left, class Right>
@@ -140,6 +140,7 @@ struct CodeProbeAnnotations<> {
 	constexpr bool expectContext(ExecutionContext context, bool prevHadSomeContext = false) const {
 		return !prevHadSomeContext;
 	}
+	constexpr bool shouldTrace(ICodeProbe const*) const { return true; }
 	constexpr bool deduplicate() const { return false; }
 };
 
@@ -161,6 +162,16 @@ struct CodeProbeAnnotations<Head, Tail...> {
 		}
 		tail.hit(self);
 	}
+
+	bool shouldTrace(ICodeProbe const* self) const {
+		if constexpr (Head::type == AnnotationType::Assertion) {
+			if (!head(self)) {
+				return false;
+			}
+		}
+		return tail.shouldTrace(self);
+	}
+
 	void trace(const ICodeProbe* self, BaseTraceEvent& evt, bool condition) const {
 		if constexpr (Head::type == AnnotationType::Decoration) {
 			head.trace(self, evt, condition);
@@ -204,6 +215,7 @@ struct ICodeProbe {
 	virtual const char* condition() const = 0;
 	virtual const char* compilationUnit() const = 0;
 	virtual void trace(bool) const = 0;
+	virtual bool shouldTrace() const = 0;
 	virtual bool wasHit() const = 0;
 	virtual unsigned hitCount() const = 0;
 	virtual bool expectInContext(ExecutionContext context) const = 0;
@@ -237,6 +249,9 @@ struct CodeProbeImpl : ICodeProbe {
 		    .detail("Comment", Comment::value());
 		annotations.trace(this, evt, condition);
 	}
+
+	bool shouldTrace() const override { return annotations.shouldTrace(this); }
+
 	bool wasHit() const override { return _hitCount > 0; }
 	unsigned hitCount() const override { return _hitCount; }