Ensure it is safe to discard unenqueued events

Summary: Similar to D85124897, we did not invalidate the handle when we override the pendingEnqueueColl_. Although this is unlikely, we should at least make it an error rather than segfault. Change the code to invalidate the handle for the old event.

Reviewed By: SuhitK

Differential Revision: D85124893

fbshipit-source-id: 7ae0696983e2fae241d46c14c1d89ef326190a86

Author: YulunW

comms/utils/colltrace/CollTrace.cc

@@ -3,6 +3,7 @@
 #include "comms/utils/colltrace/CollTrace.h"
 
 #include <fmt/core.h>
+#include <folly/json.h>
 #include <folly/logging/xlog.h>
 #include <folly/stop_watch.h>
 
@@ -82,6 +83,32 @@ CollTrace::~CollTrace() {
 CommsMaybe<std::shared_ptr<CollTraceHandle>> CollTrace::recordCollective(
     std::unique_ptr<ICollMetadata> metadata,
     std::unique_ptr<ICollWaitEvent> waitEvent) noexcept {
+  if (metadata == nullptr) {
+    return folly::makeUnexpected(CommsError(
+        "Received nullptr for metadata during recordCollective",
+        commInternalError));
+  }
+  if (waitEvent == nullptr) {
+    return folly::makeUnexpected(CommsError(
+        "Received nullptr for waitEvent during recordCollective",
+        commInternalError));
+  }
+  if (pendingEnqueueColl_ != nullptr) {
+    XLOG_FIRST_N(
+        ERR,
+        1,
+        fmt::format(
+            "{}: Got another collective enqueued when a previous one haven't finished, colltrace result would be inaccurate. Previous: {}, Next:{}",
+            logPrefix_,
+            folly::toJson(pendingEnqueueColl_->collRecord->toDynamic()),
+            folly::toJson(metadata->toDynamic())));
+    auto handlePtr = eventToHandleMap_.find(pendingEnqueueColl_.get());
+    if (handlePtr != eventToHandleMap_.end()) {
+      handlePtr->second->invalidate();
+      eventToHandleMap_.erase(pendingEnqueueColl_.get());
+    }
+  }
+
   pendingEnqueueColl_ = std::make_unique<CollTraceEvent>(
       std::make_shared<CollRecord>(collId_.fetch_add(1), std::move(metadata)),
       std::move(waitEvent));

comms/utils/colltrace/tests/CollTraceUT.cc

@@ -411,3 +411,49 @@ TEST_F(CollTraceTest, CheckHandleValidityWhenPendingQueueFull) {
     }
   }
 }
+
+// If multiple enqueue happened at the same time, colltrace would not be able
+// to handle them as there is only one slot for pending enqueue collectives.
+// We will change the behavior later to make sure we can enqueue from multiple
+// places, but for now let's at least make sure it won't cause segfault.
+TEST_F(CollTraceTest, CheckHandleValidityOverMultipleEnqueues) {
+  std::vector<std::shared_ptr<ICollTraceHandle>> handles;
+  for (int i = 0; i < 10; ++i) {
+    // Create metadata and wait event. We will only call recordCollective, so
+    // we don't expect getting any calls to them
+    auto metadata = std::make_unique<NiceMock<MockCollMetadata>>();
+    auto waitEvent = std::make_unique<NiceMock<MockCollWaitEvent>>();
+
+    // Set up expectations for the wait event
+    ON_CALL(*waitEvent, beforeCollKernelScheduled())
+        .WillByDefault(Return(folly::unit));
+
+    // Set up expectations for the wait event
+    ON_CALL(*metadata, toDynamic())
+        .WillByDefault(
+            Return(static_cast<folly::dynamic>(folly::dynamic::object())));
+
+    // Record a collective
+    auto handleMaybe =
+        collTrace->recordCollective(std::move(metadata), std::move(waitEvent));
+
+    // Verify that the handle was created successfully
+    EXPECT_VALUE(handleMaybe);
+    EXPECT_NE(handleMaybe.value().get(), nullptr);
+
+    // Trigger the enqueue
+    handleMaybe.value()->trigger(
+        CollTraceHandleTriggerState::BeforeEnqueueKernel);
+
+    handles.emplace_back(std::move(handleMaybe.value()));
+  }
+
+  for (auto& handle : handles) {
+    // Make sure we can get the coll record without encountering segmentation
+    // fault. Getting invalid record is expected.
+    auto res = handle->getCollRecord();
+    if (res.hasValue()) {
+      EXPECT_NE(res.value(), nullptr);
+    }
+  }
+}