[GitHub] [cassandra] maedhroz commented on a diff in pull request #1663: CASSANDRA-17244: Fix org.apache.cassandra.distributed.test.trackwarnings.TombstoneCountWarningTest.failThresholdSinglePartition

[GitBox <gi...@apache.org>]

maedhroz commented on code in PR #1663:
URL: https://github.com/apache/cassandra/pull/1663#discussion_r889201338


##########
test/distributed/org/apache/cassandra/distributed/test/thresholds/TombstoneCountWarningTest.java:
##########
@@ -327,4 +361,78 @@ private static ResultSet driverQueryAll(String cql)
     {
         return JAVA_DRIVER_SESSION.execute(new SimpleStatement(cql).setConsistencyLevel(com.datastax.driver.core.ConsistencyLevel.ALL));
     }
+
+    @Shared
+    public static class State
+    {
+        // use InetSocketAddress as InetAddressAndPort is @Isolated which means equality doesn't work due to different
+        // ClassLoaders; InetSocketAddress is @Shared so safe to use between app and cluster class loaders
+        public static volatile InetSocketAddress blockFor = null;
+        public static volatile CompletableFuture<Void> promise = new CompletableFuture<>();
+
+        // called on main thread
+        public static void blockFor(InetSocketAddress address)
+        {
+            blockFor = address;
+            promise = new CompletableFuture<>();
+        }
+
+        // called in C* threads; non-test threads
+        public static void onFailure(InetSocketAddress address)
+        {
+            if (address.equals(blockFor))
+                promise.complete(null);
+        }
+
+        // called on main thread
+        public static void syncAndClear()
+        {
+            if (blockFor != null)
+            {
+                promise.join();
+                blockFor = null;
+            }
+        }
+    }
+
+    public static class BB
+    {
+        private static void install(ClassLoader cl, int instanceId)
+        {
+            if (instanceId != 1)
+                return;
+            new ByteBuddy().rebase(ReadCallback.class)
+                           .method(named("awaitResults"))
+                           .intercept(MethodDelegation.to(BB.class))
+                           .method(named("onFailure"))
+                           .intercept(MethodDelegation.to(BB.class))
+                           .make()
+                           .load(cl, ClassLoadingStrategy.Default.INJECTION);
+            new ByteBuddy().rebase(SEPExecutor.class)
+                           .method(named("maybeExecuteImmediately"))
+                           .intercept(MethodDelegation.to(BB.class))
+                           .make()
+                           .load(cl, ClassLoadingStrategy.Default.INJECTION);
+        }
+
+        public static void awaitResults(@SuperCall Runnable zuper)
+        {
+            State.syncAndClear();
+            zuper.run();
+        }
+
+        public static void onFailure(InetAddressAndPort from, RequestFailureReason failureReason, @SuperCall Runnable zuper) throws Exception
+        {
+            State.onFailure(new InetSocketAddress(from.getAddress(), from.getPort()));
+            zuper.run();
+        }
+
+        // make sure to schedule the task rather than running inline...
+        // this is imporant as the read may block on the local version which can get the test to include it rather than
+        // block waiting, so by scheduling we make sure its always fair
+        public static void maybeExecuteImmediately(Runnable task, @This SEPExecutor executor)

Review Comment:
   nit: throw an `@SuppressWarnings("unused")` on this method if you're OCD like me about IDEA's compile warnings :D



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: pr-unsubscribe@cassandra.apache.org
For additional commands, e-mail: pr-help@cassandra.apache.org