Merge pull request #488 from memfork/ref-same-coroutine-id-in-same-re…

…quest Fix: ref different coroutine id on same request
swooletw · May 12, 2021 · e96b46f · e96b46f
2 parents ebc5859 + d089754
commit e96b46f
Show file tree

Hide file tree

Showing 2 changed files with 68 additions and 10 deletions.
diff --git a/src/Coroutine/Context.php b/src/Coroutine/Context.php
@@ -7,6 +7,8 @@
 
 class Context
 {
+    protected const MAX_RECURSE_COROUTINE_ID = 50;
+
     /**
      * The app containers in different coroutine environment.
      *
@@ -26,7 +28,7 @@ class Context
      */
     public static function getApp()
     {
-        return static::$apps[static::getCoroutineId()] ?? null;
+        return static::$apps[static::getRequestedCoroutineId()] ?? null;
     }
 
     /**
@@ -36,7 +38,7 @@ public static function getApp()
      */
     public static function setApp(Container $app)
     {
-        static::$apps[static::getCoroutineId()] = $app;
+        static::$apps[static::getRequestedCoroutineId()] = $app;
     }
 
     /**
@@ -48,7 +50,7 @@ public static function setApp(Container $app)
      */
     public static function getData(string $key)
     {
-        return static::$data[static::getCoroutineId()][$key] ?? null;
+        return static::$data[static::getRequestedCoroutineId()][$key] ?? null;
     }
 
     /**
@@ -59,7 +61,7 @@ public static function getData(string $key)
      */
     public static function setData(string $key, $value)
     {
-        static::$data[static::getCoroutineId()][$key] = $value;
+        static::$data[static::getRequestedCoroutineId()][$key] = $value;
     }
 
     /**
@@ -69,31 +71,46 @@ public static function setData(string $key, $value)
      */
     public static function removeData(string $key)
     {
-        unset(static::$data[static::getCoroutineId()][$key]);
+        unset(static::$data[static::getRequestedCoroutineId()][$key]);
     }
 
     /**
      * Get data keys by current coroutine id.
      */
     public static function getDataKeys()
     {
-        return array_keys(static::$data[static::getCoroutineId()] ?? []);
+        return array_keys(static::$data[static::getRequestedCoroutineId()] ?? []);
     }
 
     /**
      * Clear data by current coroutine id.
      */
     public static function clear()
     {
-        unset(static::$apps[static::getCoroutineId()]);
-        unset(static::$data[static::getCoroutineId()]);
+        unset(static::$apps[static::getRequestedCoroutineId()]);
+        unset(static::$data[static::getRequestedCoroutineId()]);
+    }
+
+    public static function getCoroutineId(): int
+    {
+        return Coroutine::getuid();
     }
 
     /**
      * Get current coroutine id.
      */
-    public static function getCoroutineId()
+    public static function getRequestedCoroutineId(): int
     {
-        return Coroutine::getuid();
+        $currentId = static::getCoroutineId();
+        if ($currentId === -1) {
+            return -1;
+        }
+
+        $counter = 0;
+        while (($topCoroutineId = Coroutine::getPcid($currentId)) !== -1 && $counter <= static::MAX_RECURSE_COROUTINE_ID) {
+            $currentId = $topCoroutineId;
+            $counter++;
+        }
+        return $currentId;
     }
 }
diff --git a/tests/Coroutine/ContextTest.php b/tests/Coroutine/ContextTest.php
@@ -52,6 +52,47 @@ public function testGetDataKeys()
         $this->assertSame(['foo', 'sea'], Context::getDataKeys());
     }
 
+    public function testGetDataKeyInCoroutine()
+    {
+        $data1 = null;
+        $data2 = null;
+        $data3 = null;
+
+        $coroutineId1 = null;
+        $coroutineId2 = null;
+        $coroutineId3 = null;
+
+        \Swoole\Coroutine\run(function () use (&$data1, &$data2, &$data3, &$coroutineId1, &$coroutineId2, &$coroutineId3) {
+            Context::setData('foo', 'bar');
+
+            $data1 = Context::getData('foo');
+            $data2 = 'baz';
+            $data2 = 'swoole';
+
+            $coroutineId1 = Context::getRequestedCoroutineId();
+            $coroutineId2 = -1;
+            $coroutineId3 = -1;
+
+            go(function () use (&$data2, &$data3, &$coroutineId2, &$coroutineId3) {
+                $data2 = Context::getData('foo');
+                $coroutineId2 = Context::getRequestedCoroutineId();
+
+                // test nested coroutine
+                go(function () use (&$data3, &$coroutineId3) {
+                    $data3 = Context::getData('foo');
+                    $coroutineId3 = Context::getRequestedCoroutineId();
+                });
+            });
+        });
+
+        $this->assertSame('bar', $data1);
+        $this->assertSame($data1, $data2);
+        $this->assertSame($data2, $data3);
+        $this->assertSame($coroutineId1, $coroutineId2);
+        $this->assertSame($coroutineId2, $coroutineId3);
+
+    }
+
     public function testClear()
     {
         Context::setApp(m::mock(Container::class));