Merge pull request #201 from /issues/200

fix(channel): reset state to EmptyChannelState after unsetting
SOF3 · May 1, 2023 · d6b666c · d6b666c
2 parents 9e07d5a + 8c61da5
commit d6b666c
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 2 deletions.
diff --git a/await-generator/src/SOFe/AwaitGenerator/Channel.php b/await-generator/src/SOFe/AwaitGenerator/Channel.php
@@ -67,7 +67,15 @@ public function sendAndWait($value) : Generator{
 			});
 		} finally {
 			if($key !== null) {
-				unset($this->state->queue[spl_object_id($key)]);
+				if($this->state instanceof SendingChannelState) {
+					// our key may still exist in the channel state
+
+					unset($this->state->queue[spl_object_id($key)]);
+					if(count($this->state->queue) === 0) {
+						$this->state = new EmptyChannelState;
+					}
+				}
+				// else, state already changed means our key has been shifted already.
 			}
 		}
 	}
@@ -134,7 +142,15 @@ public function receive() : Generator{
 			});
 		} finally {
 			if($key !== null) {
-				unset($this->state->queue[spl_object_id($key)]);
+				if($this->state instanceof ReceivingChannelState) {
+					// our key may still exist in the channel state
+
+					unset($this->state->queue[spl_object_id($key)]);
+					if(count($this->state->queue) === 0) {
+						$this->state = new EmptyChannelState;
+					}
+				}
+				// else, state already changed means our key has been shifted already.
 			}
 		}
 	}

diff --git a/tests/SOFe/AwaitGenerator/ChannelTest.php b/tests/SOFe/AwaitGenerator/ChannelTest.php
@@ -331,4 +331,43 @@ public function testTryReceive() : void{
 
 		self::assertSame("a", $receive);
 	}
+
+	public function testTryCancelSender() : void{
+		$ok = false;
+		Await::f2c(function() use(&$ok){
+			/** @var Channel<null> $channel */
+			$channel = new Channel;
+
+			[$which, $_] = yield from Await::safeRace([
+				$channel->sendAndWait(null),
+				GeneratorUtil::empty(null),
+			]);
+			self::assertSame(1, $which);
+
+			$ret = $channel->tryReceiveOr("no sender");
+			self::assertSame("no sender", $ret);
+			$ok = true;
+		});
+
+		self::assertTrue($ok, "test run complete");
+	}
+
+	public function testTryCancelReceiver() : void{
+		$ok = false;
+		Await::f2c(function() use(&$ok){
+			/** @var Channel<null> $channel */
+			$channel = new Channel;
+
+			[$which, $_] = yield from Await::safeRace([
+				$channel->receive(),
+				GeneratorUtil::empty(null),
+			]);
+			self::assertSame(1, $which);
+
+			$channel->sendWithoutWait(null);
+			$ok = true;
+		});
+
+		self::assertTrue($ok, "test run complete");
+	}
 }