[PECO-1261] Improve CloudFetch downloading queue (#209)

[PECO-1261] Wait only for the currently needed link, download the rest in background

Signed-off-by: Levko Kravets <levko.ne@gmail.com>

Author: kravets-levko

lib/result/CloudFetchResultHandler.ts

@@ -11,15 +11,15 @@ export default class CloudFetchResultHandler implements IResultsProvider<Array<B
 
   private pendingLinks: Array<TSparkArrowResultLink> = [];
 
-  private downloadedBatches: Array<Buffer> = [];
+  private downloadTasks: Array<Promise<Buffer>> = [];
 
   constructor(context: IClientContext, source: IResultsProvider<TRowSet | undefined>) {
     this.context = context;
     this.source = source;
   }
 
   public async hasMore() {
-    if (this.pendingLinks.length > 0 || this.downloadedBatches.length > 0) {
+    if (this.pendingLinks.length > 0 || this.downloadTasks.length > 0) {
       return true;
     }
     return this.source.hasMore();
@@ -32,15 +32,17 @@ export default class CloudFetchResultHandler implements IResultsProvider<Array<B
       this.pendingLinks.push(link);
     });
 
-    if (this.downloadedBatches.length === 0) {
-      const clientConfig = this.context.getConfig();
-      const links = this.pendingLinks.splice(0, clientConfig.cloudFetchConcurrentDownloads);
+    const clientConfig = this.context.getConfig();
+    const freeTaskSlotsCount = clientConfig.cloudFetchConcurrentDownloads - this.downloadTasks.length;
+
+    if (freeTaskSlotsCount > 0) {
+      const links = this.pendingLinks.splice(0, freeTaskSlotsCount);
       const tasks = links.map((link) => this.downloadLink(link));
-      const batches = await Promise.all(tasks);
-      this.downloadedBatches.push(...batches);
+      this.downloadTasks.push(...tasks);
     }
 
-    return this.downloadedBatches.splice(0, 1);
+    const batch = await this.downloadTasks.shift();
+    return batch ? [batch] : [];
   }
 
   private async downloadLink(link: TSparkArrowResultLink): Promise<Buffer> {

tests/e2e/cloudfetch.test.js

@@ -64,7 +64,7 @@ describe('CloudFetch', () => {
     // result handler should download 5 of them and schedule the rest
     expect(await cfResultHandler.hasMore()).to.be.false;
     expect(cfResultHandler.pendingLinks.length).to.be.equal(0);
-    expect(cfResultHandler.downloadedBatches.length).to.be.equal(0);
+    expect(cfResultHandler.downloadTasks.length).to.be.equal(0);
 
     sinon.spy(operation._data, 'fetchNext');
 
@@ -76,7 +76,7 @@ describe('CloudFetch', () => {
     expect(await cfResultHandler.hasMore()).to.be.true;
     // expected batches minus first 5 already fetched
     expect(cfResultHandler.pendingLinks.length).to.be.equal(resultLinksCount - cloudFetchConcurrentDownloads);
-    expect(cfResultHandler.downloadedBatches.length).to.be.equal(cloudFetchConcurrentDownloads - 1);
+    expect(cfResultHandler.downloadTasks.length).to.be.equal(cloudFetchConcurrentDownloads - 1);
 
     let fetchedRowCount = chunk.length;
     while (await operation.hasMoreRows()) {

tests/unit/result/CloudFetchResultHandler.test.js

@@ -89,19 +89,19 @@ describe('CloudFetchResultHandler', () => {
 
     case1: {
       result.pendingLinks = [];
-      result.downloadedBatches = [];
+      result.downloadTasks = [];
       expect(await result.hasMore()).to.be.false;
     }
 
     case2: {
       result.pendingLinks = [{}]; // just anything here
-      result.downloadedBatches = [];
+      result.downloadTasks = [];
       expect(await result.hasMore()).to.be.true;
     }
 
     case3: {
       result.pendingLinks = [];
-      result.downloadedBatches = [{}]; // just anything here
+      result.downloadTasks = [{}]; // just anything here
       expect(await result.hasMore()).to.be.true;
     }
   });
@@ -134,19 +134,19 @@ describe('CloudFetchResultHandler', () => {
     } while (await rowSetProvider.hasMore());
 
     expect(result.pendingLinks.length).to.be.equal(expectedLinksCount);
-    expect(result.downloadedBatches.length).to.be.equal(0);
+    expect(result.downloadTasks.length).to.be.equal(0);
     expect(result.fetch.called).to.be.false;
   });
 
   it('should download batches according to settings', async () => {
     const clientConfig = DBSQLClient.getDefaultConfig();
-    clientConfig.cloudFetchConcurrentDownloads = 2;
+    clientConfig.cloudFetchConcurrentDownloads = 3;
 
     const rowSet = {
       startRowOffset: 0,
       resultLinks: [...sampleRowSet1.resultLinks, ...sampleRowSet2.resultLinks],
     };
-    const expectedLinksCount = rowSet.resultLinks.length;
+    const expectedLinksCount = rowSet.resultLinks.length; // 5
     const rowSetProvider = new ResultsProviderMock([rowSet]);
     const context = {
       getConfig: () => clientConfig,
@@ -166,24 +166,28 @@ describe('CloudFetchResultHandler', () => {
     expect(await rowSetProvider.hasMore()).to.be.true;
 
     initialFetch: {
+      // `cloudFetchConcurrentDownloads` out of `expectedLinksCount` links should be scheduled immediately
+      // first one should be `await`-ed and returned from `fetchNext`
       const items = await result.fetchNext({ limit: 10000 });
       expect(items.length).to.be.gt(0);
       expect(await rowSetProvider.hasMore()).to.be.false;
 
       expect(result.fetch.callCount).to.be.equal(clientConfig.cloudFetchConcurrentDownloads);
       expect(result.pendingLinks.length).to.be.equal(expectedLinksCount - clientConfig.cloudFetchConcurrentDownloads);
-      expect(result.downloadedBatches.length).to.be.equal(clientConfig.cloudFetchConcurrentDownloads - 1);
+      expect(result.downloadTasks.length).to.be.equal(clientConfig.cloudFetchConcurrentDownloads - 1);
     }
 
     secondFetch: {
-      // It should return previously fetched batch, not performing additional network requests
+      // It should return previously fetched batch, and schedule one more
       const items = await result.fetchNext({ limit: 10000 });
       expect(items.length).to.be.gt(0);
       expect(await rowSetProvider.hasMore()).to.be.false;
 
-      expect(result.fetch.callCount).to.be.equal(clientConfig.cloudFetchConcurrentDownloads); // no new fetches
-      expect(result.pendingLinks.length).to.be.equal(expectedLinksCount - clientConfig.cloudFetchConcurrentDownloads);
-      expect(result.downloadedBatches.length).to.be.equal(clientConfig.cloudFetchConcurrentDownloads - 2);
+      expect(result.fetch.callCount).to.be.equal(clientConfig.cloudFetchConcurrentDownloads + 1);
+      expect(result.pendingLinks.length).to.be.equal(
+        expectedLinksCount - clientConfig.cloudFetchConcurrentDownloads - 1,
+      );
+      expect(result.downloadTasks.length).to.be.equal(clientConfig.cloudFetchConcurrentDownloads - 1);
     }
 
     thirdFetch: {
@@ -192,11 +196,11 @@ describe('CloudFetchResultHandler', () => {
       expect(items.length).to.be.gt(0);
       expect(await rowSetProvider.hasMore()).to.be.false;
 
-      expect(result.fetch.callCount).to.be.equal(clientConfig.cloudFetchConcurrentDownloads * 2);
+      expect(result.fetch.callCount).to.be.equal(clientConfig.cloudFetchConcurrentDownloads + 2);
       expect(result.pendingLinks.length).to.be.equal(
-        expectedLinksCount - clientConfig.cloudFetchConcurrentDownloads * 2,
+        expectedLinksCount - clientConfig.cloudFetchConcurrentDownloads - 2,
       );
-      expect(result.downloadedBatches.length).to.be.equal(clientConfig.cloudFetchConcurrentDownloads - 1);
+      expect(result.downloadTasks.length).to.be.equal(clientConfig.cloudFetchConcurrentDownloads - 1);
     }
   });
 
@@ -251,6 +255,10 @@ describe('CloudFetchResultHandler', () => {
       }),
     );
 
+    // There are two link in the batch - first one is valid and second one is expired
+    // The first fetch has to be successful, and the second one should fail
+    await result.fetchNext({ limit: 10000 });
+
     try {
       await result.fetchNext({ limit: 10000 });
       expect.fail('It should throw an error');