[BUG]Test simple pinned blocking alloc Failed nightly tests

[tgravescs]

Describe the bug
rapids_nightly-dev-github #1073 failed due to spark 334 unit test - simple pinned blocking alloc failing. Looks like it timed out.

[2024-03-13T13:29:55.623Z]              TEST THREAD APPEARS TO BE STUCK
[2024-03-13T13:29:55.623Z] thread2
[2024-03-13T13:29:55.624Z]      com.nvidia.spark.rapids.jni.SparkResourceAdaptor.preCpuAlloc(Native Method)
[2024-03-13T13:29:55.624Z]      com.nvidia.spark.rapids.jni.SparkResourceAdaptor.preCpuAlloc(SparkResourceAdaptor.java:262)
[2024-03-13T13:29:55.624Z]      com.nvidia.spark.rapids.jni.RmmSpark.preCpuAlloc(RmmSpark.java:607)
[2024-03-13T13:29:55.624Z]      com.nvidia.spark.rapids.HostAlloc.tryAllocInternal(HostAlloc.scala:162)
[2024-03-13T13:29:55.624Z]      com.nvidia.spark.rapids.HostAlloc.alloc(HostAlloc.scala:218)
[2024-03-13T13:29:55.624Z]      com.nvidia.spark.rapids.HostAlloc$.alloc(HostAlloc.scala:258)
[2024-03-13T13:29:55.624Z]      com.nvidia.spark.rapids.HostAllocSuite$AllocOnAnotherThread.$anonfun$doAlloc$1(HostAllocSuite.scala:269)
[2024-03-13T13:29:55.624Z]      com.nvidia.spark.rapids.HostAllocSuite$AllocOnAnotherThread$$Lambda$570/563871766.apply$mcV$sp(Unknown Source)
[2024-03-13T13:29:55.624Z]      scala.runtime.java8.JFunction0$mcV$sp.apply(JFunction0$mcV$sp.java:23)

[2024-03-13T13:29:55.624Z] TEST THREAD ScalaTest-main-running-HostAllocSuite
[2024-03-13T13:29:55.624Z]      java.lang.Object.wait(Native Method)
[2024-03-13T13:29:55.624Z]      com.nvidia.spark.rapids.HostAllocSuite$TaskThread$TaskThreadTrackingOp.get(HostAllocSuite.scala:111)
[2024-03-13T13:29:55.624Z]      com.nvidia.spark.rapids.HostAllocSuite$AllocOnAnotherThread.waitForAlloc(HostAllocSuite.scala:221)
[2024-03-13T13:29:55.624Z]      com.nvidia.spark.rapids.HostAllocSuite.$anonfun$new$24(HostAllocSuite.scala:454)
[2024-03-13T13:29:55.624Z]      com.nvidia.spark.rapids.HostAllocSuite.$anonfun$new$24$adapted(HostAllocSuite.scala:449)
[2024-03-13T13:29:55.624Z]      com.nvidia.spark.rapids.HostAllocSuite$$Lambda$580/1058189103.apply(Unknown Source)
[2024-03-13T13:29:55.624Z]      com.nvidia.spark.rapids.Arm$.withResource(Arm.scala:30)
[2024-03-13T13:29:55.624Z]      com.nvidia.spark.rapids.HostAllocSuite.$anonfun$new$23(HostAllocSuite.scala:449)
[2024-03-13T13:29:55.624Z]      com.nvidia.spark.rapids.HostAllocSuite.$anonfun$new$23$adapted(HostAllocSuite.scala:445)
[2024-03-13T13:29:55.624Z]      com.nvidia.spark.rapids.HostAllocSuite$$Lambda$579/809665906.apply(Unknown Source)
[2024-03-13T13:29:55.624Z]      com.nvidia.spark.rapids.Arm$.withResource(Arm.scala:30)
[2024-03-13T13:29:55.624Z]      com.nvidia.spark.rapids.HostAllocSuite.$anonfun$new$21(HostAllocSuite.scala:445)
[2024-03-13T13:29:55.624Z]      com.nvidia.spark.rapids.HostAllocSuite$$Lambda$482/860099309.apply(Unknown Source)

[abellina]

I can repro this and I am looking at the issue. It appears to be a race condition where a free/notify is not seen by the state machine but I don't know why. It repros pretty consistently, and adding transition logging makes it go away.

[abellina]

Ok here's what I see so far, the state in HostAlloc and the state machine in spark-rapids-jni is getting updated via MemoryBuffer.onClosed. This callback is called before we free the memory from the pinned pool. Because of that we update state in HostAlloc saying that we have free pinned memory and we should attempt an allocation, but what ends up happening is that we fail to allocate because the pinned pool hasn't seen the free yet (onClosed triggered, but we haven't gotten to the free yet).

In this case then HostAlloc will NOT retry the allocation and instead it sends the failing thread to BLOCKED state. Unfortunately the event that would have moved this thread out of BLOCKED already triggered via onClosed so we "miss the free" in HostAlloc.