SNOW-734920: Add possibility to configure keep alive value

snowflakedb · Sep 21, 2023 · 9e1faa8 · 9e1faa8
1 parent a6367ff
commit 9e1faa8
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 24 deletions.
diff --git a/lib/http/node.js b/lib/http/node.js
@@ -44,11 +44,10 @@ Util.inherits(NodeHttpClient, Base);
 
 const httpsAgentCache = new Map();
 
-function getFromCacheOrCreate(agentClass, options, url) {
-  const parsed = Url.parse(url);
-  const protocol = parsed.protocol || 'http:';
-  const port = parsed.port || (protocol === 'http:' ? '80' : '443');
-  const agentId = `${protocol}//${parsed.hostname}:${port}/${options.keepAlive}`
+function getFromCacheOrCreate(agentClass, options, parsedUrl) {
+  const protocol = parsedUrl.protocol || 'http:';
+  const port = parsedUrl.port || (protocol === 'http:' ? '80' : '443');
+  const agentId = `${protocol}//${parsedUrl.hostname}:${port}/${options.keepAlive}`
   if (httpsAgentCache.has(agentId)) {
     Logger.getInstance().trace(`Get agent with id: ${agentId} from cache`);
     return httpsAgentCache.get(agentId);
@@ -61,7 +60,6 @@ function getFromCacheOrCreate(agentClass, options, url) {
 }
 
 function prepareProxyAgentOptions(agentOptions, proxy) {
-  Logger.getInstance().info(`Use proxy: ${JSON.stringify(proxy)}`);
   agentOptions.host = proxy.host;
   agentOptions.port = proxy.port;
   agentOptions.protocol = proxy.protocol;
@@ -90,28 +88,31 @@ function isBypassProxy(proxy, url, bypassProxy) {
 /**
  * @inheritDoc
  */
-NodeHttpClient.prototype.getAgent = function (url, proxy, mock, defaultAgentOptions) {
-  const isHttps = Url.parse(url).protocol === 'https:';
-  let agentOptions = {keepAlive: GlobalConfig.keepAlive()};
-  var bypassProxy = isBypassProxy(proxy, url);
-
+NodeHttpClient.prototype.getAgent = function (url, proxy, mock) {
+  const agentOptions = {keepAlive: GlobalConfig.keepAlive()};
   if (mock) {
-    return mock.agentClass(agentOptions)
+    return mock.agentClass(agentOptions);
   }
 
+  const parsedUrl = Url.parse(url);
+  const isHttps = parsedUrl.protocol === 'https:';
+  const bypassProxy = isBypassProxy(proxy, url);
+  let agent = {};
+
   if (isHttps) {
     if (proxy && !bypassProxy) {
-      const proxyAgentOptions = prepareProxyAgentOptions(agentOptions, proxy)
-      return getFromCacheOrCreate(HttpsProxyAgent, proxyAgentOptions, url);
+      const proxyAgentOptions = prepareProxyAgentOptions(agentOptions, proxy);
+      agent = getFromCacheOrCreate(HttpsProxyAgent, proxyAgentOptions, parsedUrl);
     } else {
-      return getFromCacheOrCreate(HttpsAgent, agentOptions, url);
+      agent = getFromCacheOrCreate(HttpsAgent, agentOptions, parsedUrl);
     }
   } else if (proxy && !bypassProxy) {
     const proxyAgentOptions = prepareProxyAgentOptions(agentOptions, proxy);
-    return getFromCacheOrCreate(HttpAgent, proxyAgentOptions, url);
+    agent = getFromCacheOrCreate(HttpAgent, proxyAgentOptions, url);
   } else {
-    return getFromCacheOrCreate(HttpAgent, agentOptions, url);
+    agent = getFromCacheOrCreate(HttpAgent, agentOptions, parsedUrl);
   }
+  return agent;
 };
 
 module.exports = NodeHttpClient;
diff --git a/test/integration/testKeepAlive.js b/test/integration/testKeepAlive.js
@@ -12,19 +12,24 @@ const assert = require("assert");
 describe('keepAlive test', function ()
 {
   let connection;
-  const loopCount = 5;
+  const loopCount = 10;
   const rowCount = 10;
   const tableName = 'test_keepalive000';
 
   const createTableWithRandomStrings = `CREATE OR REPLACE TABLE ${tableName} (value string)
     AS select randstr(200, random()) from table (generator(rowcount =>${rowCount}))`;
 
+  before(async () => {
+    connection = snowflake.createConnection(connOption.valid);
+    await testUtil.connectAsync(connection);
+    await testUtil.executeCmdAsync(connection, createTableWithRandomStrings);
+  });
   after(async () => {
     await testUtil.dropTablesIgnoringErrorsAsync(connection, [tableName]);
     await testUtil.destroyConnectionAsync(connection);
   });
 
-  it('Run query in loop', function (done)
+  it('Verify that requests working faster with set keep alive', function (done)
   {
     let sumWithKeepAlive = 0;
     let sumWithoutKeepAlive = 0;
@@ -33,9 +38,6 @@ describe('keepAlive test', function ()
         // Run the query loopCount times
         async function ()
         {
-          connection = snowflake.createConnection(connOption.valid);
-          await testUtil.connectAsync(connection);
-          await testUtil.executeCmdAsync(connection, createTableWithRandomStrings);
           for (let count = 1; count <= loopCount; count++)
           {
             await new Promise((resolve, reject) =>
@@ -72,7 +74,6 @@ describe('keepAlive test', function ()
           snowflake.configure({keepAlive: false});
           connection = snowflake.createConnection(connOption.valid);
           await testUtil.connectAsync(connection);
-          await testUtil.executeCmdAsync(connection, createTableWithRandomStrings);
           for (let count = 1; count <= loopCount; count++)
           {
             await new Promise((resolve, reject) =>
@@ -106,7 +107,7 @@ describe('keepAlive test', function ()
         },
         async function ()
         {
-          assert.ok(sumWithoutKeepAlive/2 > sumWithKeepAlive, 'With keep alive queries should work more thdn two times faster');
+          assert.ok(sumWithoutKeepAlive/2 > sumWithKeepAlive, 'With keep alive queries should work more than two times faster');
         }
       ],
       done