Support custom pool key (#127)

amuluowin · twose · web-flow · commit 5c240e71aaa4 · 2021-05-18T09:49:15.000+08:00
Co-authored-by: twosee &lt;twose@qq.com&gt;
diff --git a/README.md b/README.md
@@ -486,6 +486,7 @@ go(function(){
 | retry                 | callable              | 自动重试拦截器     | `function(Request $request, Response $response){}`           | 位于发生错误后及重试之前                                     |
 | retry_time            | int                   | 自动重试次数       |                                                              | 默认不重试                                                   |
 | use_pool              | bool\|int             | 连接池             | `true`                                                       | `false`                                                      |
+| pool_key              | callable\|array       | 连接池的key        | `function(Request $request):string { return $key; }`         | 默认为请求地址的`host:port`                                  |
 
 ### 配置参数别名
 
diff --git a/src/ClientPool.php b/src/ClientPool.php
@@ -25,18 +25,18 @@ public function createEx(array $options, bool $temp = false)
             );
         }
         $client = new Client($options['host'], $options['port'], $options['ssl']);
-        if ($temp) {
-            return $client; // not record
-        } else {
-            $key = "{$options['host']}:{$options['port']}";
+        if (!$temp) {
+            $key = $options['pool_key'] ?? "{$options['host']}:{$options['port']}";
             parent::create($options, $key);
-            return $client;
+            /** @noinspection PhpUndefinedFieldInspection */
+            $client->pool_key = $key;
         }
+        return $client;
     }
 
     public function setMaxEx(array $options, int $max_size = -1): int
     {
-        $key = "{$options['host']}:{$options['port']}";
+        $key = $options['pool_key'] ?? "{$options['host']}:{$options['port']}";
         $ret = parent::setMax($key, $max_size);
         if ($ret === -1) { // chan reduce max size
             $chan = $this->resource_map[$key];
@@ -54,10 +54,10 @@ public function setMaxEx(array $options, int $max_size = -1): int
         return $ret;
     }
 
-    public function getEx(string $host, string $port): ?Client
+    public function getEx(array $options): ?Client
     {
         /** @var $client Client */
-        $key = "{$host}:{$port}";
+        $key = $options['pool_key'] ?? "{$options['host']}:{$options['port']}";
         $client = parent::get($key);
         if ($client && SABER_SW_LE_V401 && !$client->isConnected()) {
             @$this->status_map[$key]['disconnected']++;
@@ -68,17 +68,13 @@ public function getEx(string $host, string $port): ?Client
 
     public function putEx(Client $client)
     {
-        /** @var $client Client */
-        if (!($client instanceof Client)) {
-            throw new InvalidArgumentException('$value should be instance of ' . Client::class);
-        }
-        parent::put($client, "{$client->host}:{$client->port}");
+        parent::put($client, $client->pool_key ?? "{$client->host}:{$client->port}");
     }
 
     public function destroyEx(Client $client)
     {
         $client->close();
-        parent::destroy($client, "{$client->host}:{$client->port}");
+        parent::destroy($client, $client->pool_key ?? "{$client->host}:{$client->port}");
     }
 
     public function release(string $key)
diff --git a/src/Request.php b/src/Request.php
@@ -215,7 +215,11 @@ public function withPool($bool_or_max_size): self
         $this->use_pool = $bool_or_max_size;
         if (is_numeric($this->use_pool)) {
             // limit max num
-            ClientPool::getInstance()->setMaxEx($this->getConnectionTarget(), $bool_or_max_size);
+            $connectionTarget = $this->getConnectionTarget();
+            if (($key = $this->callInterceptor('pool_key', $this))) {
+                $connectionTarget['pool_key'] = $key;
+            }
+            ClientPool::getInstance()->setMaxEx($connectionTarget, $bool_or_max_size);
         }
         if ($bool_or_max_size) {
             $this->withKeepAlive(true);
@@ -650,11 +654,13 @@ public function exec()
             $connectionInfo = $this->getConnectionTarget();
             /** create a new coroutine client */
             $client_pool = ClientPool::getInstance();
-            if ($this->use_pool && $client = $client_pool->getEx($connectionInfo['host'], $connectionInfo['port'])) {
-                $this->client = $client;
-            } else {
-                $this->client = $client_pool->createEx($connectionInfo, !$this->use_pool);
+            if ($this->use_pool) {
+                if (($key = $this->callInterceptor('pool_key', $this))) {
+                    $connectionInfo['pool_key'] = $key;
+                }
+                $client = $client_pool->getEx($connectionInfo);
             }
+            $this->client = $client ?? $client_pool->createEx($connectionInfo, !$this->use_pool);
         }
 
         /** Clear useless cookies property */
diff --git a/src/Saber.php b/src/Saber.php
@@ -469,10 +469,6 @@ private static function transOptionsToRequest(array $options, Request $request)
             );
         }
 
-        if (isset($options['use_pool'])) {
-            $request->withPool($options['use_pool']);
-        }
-
         if (!empty($options['uri_query']) && $uri = $request->getUri()) {
             $uri->withQuery(DataParser::toQueryString(($options['uri_query'])));
         }
@@ -661,6 +657,10 @@ private static function transOptionsToRequest(array $options, Request $request)
             $request->withAddedInterceptor('before_retry', (array)$options['before_retry']);
         }
 
+        if(!empty($options['pool_key'])){
+            $request->withAddedInterceptor('pool_key', (array)$options['pool_key']);
+        }
+
         /** register Interceptor before request */
         if (!empty($options['before'])) {
             $request->withAddedInterceptor('request', (array)$options['before']);
@@ -679,6 +679,10 @@ private static function transOptionsToRequest(array $options, Request $request)
         if (!empty($options['exception_handle'])) {
             $request->withAddedInterceptor('exception', (array)$options['exception_handle']);
         }
+
+        if (isset($options['use_pool'])) {
+            $request->withPool($options['use_pool']);
+        }
     }
 
     private static function mergeOptions(array $options, ... $defaults)
diff --git a/tests/SaberTest.php b/tests/SaberTest.php
@@ -427,4 +427,18 @@ public function testNonRedirect()
         $res = $saber->get('http://baidu.com', ['redirect' => 0]);
         $this->assertTrue((string)$res->body !== '');
     }
+
+    public function testPoolKeyConfig(){
+        $saber = Saber::create(['exception_report' => true]);
+        $saber->get('http://www.httpbin.org', ['use_pool' => 5, 'pool_key' => function (Saber\Request $request) {
+            return 'test_pool';
+        }]);
+        $this->assertEquals(saber_pool_get_status('test_pool')['max'], 5);
+    }
+
+    public function testNoPoolKeyConfig(){
+        $saber = Saber::create(['exception_report' => true]);
+        $saber->get('http://www.httpbin.org', ['use_pool' => 5]);
+        $this->assertEquals(saber_pool_get_status('www.httpbin.org:80')['max'], 5);
+    }
 }

Original file line number	Diff line number	Diff line change
`@@ -25,18 +25,18 @@ public function createEx(array $options, bool $temp = false)`
`25`	`25`	`);`
`26`	`26`	`}`
`27`	`27`	`$client = new Client($options['host'], $options['port'], $options['ssl']);`
`28`		`- if ($temp) {`
`29`		`- return $client; // not record`
`30`		`- } else {`
`31`		`- $key = "{$options['host']}:{$options['port']}";`
	`28`	`+ if (!$temp) {`
	`29`	`+ $key = $options['pool_key'] ?? "{$options['host']}:{$options['port']}";`
`32`	`30`	`parent::create($options, $key);`
`33`		`- return $client;`
	`31`	`+ /** @noinspection PhpUndefinedFieldInspection */`
	`32`	`+ $client->pool_key = $key;`
`34`	`33`	`}`
	`34`	`+ return $client;`
`35`	`35`	`}`
`36`	`36`
`37`	`37`	`public function setMaxEx(array $options, int $max_size = -1): int`
`38`	`38`	`{`
`39`		`- $key = "{$options['host']}:{$options['port']}";`
	`39`	`+ $key = $options['pool_key'] ?? "{$options['host']}:{$options['port']}";`
`40`	`40`	`$ret = parent::setMax($key, $max_size);`
`41`	`41`	`if ($ret === -1) { // chan reduce max size`
`42`	`42`	`$chan = $this->resource_map[$key];`
`@@ -54,10 +54,10 @@ public function setMaxEx(array $options, int $max_size = -1): int`
`54`	`54`	`return $ret;`
`55`	`55`	`}`
`56`	`56`
`57`		`- public function getEx(string $host, string $port): ?Client`
	`57`	`+ public function getEx(array $options): ?Client`
`58`	`58`	`{`
`59`	`59`	`/** @var $client Client */`
`60`		`- $key = "{$host}:{$port}";`
	`60`	`+ $key = $options['pool_key'] ?? "{$options['host']}:{$options['port']}";`
`61`	`61`	`$client = parent::get($key);`
`62`	`62`	`if ($client && SABER_SW_LE_V401 && !$client->isConnected()) {`
`63`	`63`	`@$this->status_map[$key]['disconnected']++;`
`@@ -68,17 +68,13 @@ public function getEx(string $host, string $port): ?Client`
`68`	`68`
`69`	`69`	`public function putEx(Client $client)`
`70`	`70`	`{`
`71`		`- /** @var $client Client */`
`72`		`- if (!($client instanceof Client)) {`
`73`		`- throw new InvalidArgumentException('$value should be instance of ' . Client::class);`
`74`		`- }`
`75`		`- parent::put($client, "{$client->host}:{$client->port}");`
	`71`	`+ parent::put($client, $client->pool_key ?? "{$client->host}:{$client->port}");`
`76`	`72`	`}`
`77`	`73`
`78`	`74`	`public function destroyEx(Client $client)`
`79`	`75`	`{`
`80`	`76`	`$client->close();`
`81`		`- parent::destroy($client, "{$client->host}:{$client->port}");`
	`77`	`+ parent::destroy($client, $client->pool_key ?? "{$client->host}:{$client->port}");`
`82`	`78`	`}`
`83`	`79`
`84`	`80`	`public function release(string $key)`
Original file line number	Diff line number	Diff line change
`@@ -469,10 +469,6 @@ private static function transOptionsToRequest(array $options, Request $request)`
`469`	`469`	`);`
`470`	`470`	`}`
`471`	`471`
`472`		`- if (isset($options['use_pool'])) {`
`473`		`- $request->withPool($options['use_pool']);`
`474`		`- }`
`475`		`-`
`476`	`472`	`if (!empty($options['uri_query']) && $uri = $request->getUri()) {`
`477`	`473`	`$uri->withQuery(DataParser::toQueryString(($options['uri_query'])));`
`478`	`474`	`}`
`@@ -661,6 +657,10 @@ private static function transOptionsToRequest(array $options, Request $request)`
`661`	`657`	`$request->withAddedInterceptor('before_retry', (array)$options['before_retry']);`
`662`	`658`	`}`
`663`	`659`
	`660`	`+ if(!empty($options['pool_key'])){`
	`661`	`+ $request->withAddedInterceptor('pool_key', (array)$options['pool_key']);`
	`662`	`+ }`
	`663`	`+`
`664`	`664`	`/** register Interceptor before request */`
`665`	`665`	`if (!empty($options['before'])) {`
`666`	`666`	`$request->withAddedInterceptor('request', (array)$options['before']);`
`@@ -679,6 +679,10 @@ private static function transOptionsToRequest(array $options, Request $request)`
`679`	`679`	`if (!empty($options['exception_handle'])) {`
`680`	`680`	`$request->withAddedInterceptor('exception', (array)$options['exception_handle']);`
`681`	`681`	`}`
	`682`	`+`
	`683`	`+ if (isset($options['use_pool'])) {`
	`684`	`+ $request->withPool($options['use_pool']);`
	`685`	`+ }`
`682`	`686`	`}`
`683`	`687`
`684`	`688`	`private static function mergeOptions(array $options, ... $defaults)`