SWDEV-393637 - Standardize the way to implement wait mechanism (#316)

Change-Id: I59027667806878191f9c641f27ce47e6f85ba40c

catch/include/hip_test_common.hh

@@ -22,6 +22,7 @@ THE SOFTWARE.
 
 #pragma once
 #include "hip_test_context.hh"
+
 #include <catch.hpp>
 #include <atomic>
 #include <chrono>
@@ -30,6 +31,7 @@ THE SOFTWARE.
 #include <iomanip>
 #include <mutex>
 #include <cstdlib>
+#include <thread>
 
 #define HIP_PRINT_STATUS(status) INFO(hipGetErrorName(status) << " at line: " << __LINE__);
 
@@ -429,6 +431,52 @@ static inline void runKernelForDuration(std::chrono::milliseconds duration,
   hipLaunchKernelGGL(waitKernel_used, dim3(1), dim3(1), 0, stream, ticksPerSecond * millis / 1000);
 }
 
+class BlockingContext {
+  std::atomic_bool blocked{true};
+  hipStream_t stream;
+
+ public:
+  BlockingContext(hipStream_t s) : stream(s), blocked(true) {}
+
+  BlockingContext(const BlockingContext& in) {
+    blocked = in.blocked_val();
+    stream = in.stream_val();
+  }
+
+  BlockingContext(const BlockingContext&& in) {
+    blocked = in.blocked_val();
+    stream = in.stream_val();
+  }
+
+  void reset() { blocked = true; }
+
+  BlockingContext& operator=(const BlockingContext& in) {
+    blocked = in.blocked_val();
+    stream = in.stream_val();
+    return *this;
+  }
+
+  void block_stream() {
+    blocked = true;
+    auto blocking_callback = [](hipStream_t, hipError_t, void* data) {
+      auto blocked = reinterpret_cast<std::atomic_bool*>(data);
+      while (blocked->load()) {
+        // Yield this thread till we are waiting
+        std::this_thread::yield();
+      }
+    };
+    HIP_CHECK(hipStreamAddCallback(stream, blocking_callback, (void*)&blocked, 0));
+  }
+
+  void unblock_stream() {
+    blocked = false;
+  }
+
+  bool is_blocked() const { return hipStreamQuery(stream) == hipErrorNotReady; }
+
+  bool blocked_val() const { return blocked.load(); }
+  hipStream_t stream_val() const { return stream; }
+};
 }  // namespace HipTest
 
 // This must be called in the beginning of image test app's main() to indicate whether image