From cfb4465c38370622ddf6d66d3f6f1a25d90c65b3 Mon Sep 17 00:00:00 2001
From: MediaPipe Team <mediapipe-team@google.com>
Date: Wed, 20 Dec 2023 18:19:52 -0800
Subject: [PATCH] Guard WaitOnGpu with extra OpenGL checks.

PiperOrigin-RevId: 592707594
---
 mediapipe/gpu/BUILD         |  1 +
 mediapipe/gpu/gl_context.cc | 33 +++++++++++++++++++++++++++++++++
 mediapipe/gpu/gl_context.h  |  5 +++++
 3 files changed, 39 insertions(+)

diff --git a/mediapipe/gpu/BUILD b/mediapipe/gpu/BUILD
index c4ddfb559b..de38598024 100644
--- a/mediapipe/gpu/BUILD
+++ b/mediapipe/gpu/BUILD
@@ -210,6 +210,7 @@ cc_library(
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/synchronization",
     ] + select({
         "//conditions:default": [],
diff --git a/mediapipe/gpu/gl_context.cc b/mediapipe/gpu/gl_context.cc
index 4fda8da439..54d330074b 100644
--- a/mediapipe/gpu/gl_context.cc
+++ b/mediapipe/gpu/gl_context.cc
@@ -26,6 +26,7 @@
 #include "absl/log/absl_log.h"
 #include "absl/memory/memory.h"
 #include "absl/status/status.h"
+#include "absl/strings/str_format.h"
 #include "absl/synchronization/mutex.h"
 #include "mediapipe/framework/port.h"  // IWYU pragma: keep
 #include "mediapipe/framework/port/ret_check.h"
@@ -650,6 +651,11 @@ class GlSyncWrapper {
     // TODO: do something if the wait fails?
   }
 
+  // This method exists only for investigation purposes to distinguish stack
+  // traces: external vs. internal context.
+  // TODO: remove after glWaitSync crashes are resolved.
+  void WaitOnGpuExternalContext() { glWaitSync(sync_, 0, GL_TIMEOUT_IGNORED); }
+
   void WaitOnGpu() {
     if (!sync_) return;
       // WebGL2 specifies a waitSync call, but since cross-context
@@ -657,6 +663,33 @@ class GlSyncWrapper {
       // a warning when it's called, so let's just skip the call. See
       // b/184637485 for details.
 #ifndef __EMSCRIPTEN__
+
+    if (!GlContext::IsAnyContextCurrent()) {
+      // glWaitSync must be called on with some context current. Doing the
+      // opposite doesn't necessarily result in a crash or GL error. Hence,
+      // just logging an error and skipping the call.
+      ABSL_LOG_FIRST_N(ERROR, 1)
+          << "An attempt to wait for a sync without any context current.";
+      return;
+    }
+
+    auto context = GlContext::GetCurrent();
+    if (context == nullptr) {
+      // This can happen when WaitOnGpu is invoked on an external context,
+      // created by other than GlContext::Create means.
+      WaitOnGpuExternalContext();
+      return;
+    }
+
+    // GlContext::ShouldUseFenceSync guards creation of sync objects, so this
+    // CHECK should never fail if clients use MediaPipe APIs in an intended way.
+    // TODO: remove after glWaitSync crashes are resolved.
+    ABSL_CHECK(context->ShouldUseFenceSync()) << absl::StrFormat(
+        "An attempt to wait for a sync when it should not be used. (OpenGL "
+        "Version "
+        "%d.%d)",
+        context->gl_major_version(), context->gl_minor_version());
+
     glWaitSync(sync_, 0, GL_TIMEOUT_IGNORED);
 #endif
   }
diff --git a/mediapipe/gpu/gl_context.h b/mediapipe/gpu/gl_context.h
index 8d69483051..19dbce2429 100644
--- a/mediapipe/gpu/gl_context.h
+++ b/mediapipe/gpu/gl_context.h
@@ -71,6 +71,8 @@ typedef std::function<void()> GlVoidFunction;
 typedef std::function<absl::Status()> GlStatusFunction;
 
 class GlContext;
+// TODO: remove after glWaitSync crashes are resolved.
+class GlSyncWrapper;
 
 // Generic interface for synchronizing access to a shared resource from a
 // different context. This is an abstract class to keep users from
@@ -329,6 +331,9 @@ class GlContext : public std::enable_shared_from_this<GlContext> {
       SyncTokenTypeForTest type);
 
  private:
+  // TODO: remove after glWaitSync crashes are resolved.
+  friend GlSyncWrapper;
+
   GlContext();
 
   bool ShouldUseFenceSync() const;