[Backport 7.x] Added timeout support in bulk and msearch helpers (#1211)

Co-authored-by: Tomas Della Vedova <delvedor@users.noreply.github.com>

test/unit/helpers/bulk.test.js

@@ -7,6 +7,7 @@
 const { createReadStream } = require('fs')
 const { join } = require('path')
 const split = require('split2')
+const FakeTimers = require('@sinonjs/fake-timers')
 const semver = require('semver')
 const { test } = require('tap')
 const { Client, errors } = require('../../../')
@@ -987,3 +988,118 @@ test('errors', t => {
 
   t.end()
 })
+
+test('Flush interval', t => {
+  t.test('Slow producer', async t => {
+    const clock = FakeTimers.install({ toFake: ['setTimeout', 'clearTimeout'] })
+    t.teardown(() => clock.uninstall())
+
+    let count = 0
+    const MockConnection = connection.buildMockConnection({
+      onRequest (params) {
+        t.strictEqual(params.path, '/_bulk')
+        t.match(params.headers, { 'content-type': 'application/x-ndjson' })
+        const [action, payload] = params.body.split('\n')
+        t.deepEqual(JSON.parse(action), { index: { _index: 'test' } })
+        t.deepEqual(JSON.parse(payload), dataset[count++])
+        return { body: { errors: false, items: [{}] } }
+      }
+    })
+
+    const client = new Client({
+      node: 'http://localhost:9200',
+      Connection: MockConnection
+    })
+
+    const result = await client.helpers.bulk({
+      datasource: (async function * generator () {
+        for (const chunk of dataset) {
+          await clock.nextAsync()
+          yield chunk
+        }
+      })(),
+      flushBytes: 5000000,
+      concurrency: 1,
+      onDocument (doc) {
+        return {
+          index: { _index: 'test' }
+        }
+      },
+      onDrop (doc) {
+        t.fail('This should never be called')
+      }
+    })
+
+    t.type(result.time, 'number')
+    t.type(result.bytes, 'number')
+    t.match(result, {
+      total: 3,
+      successful: 3,
+      retry: 0,
+      failed: 0,
+      aborted: false
+    })
+  })
+
+  t.test('Abort operation', async t => {
+    const clock = FakeTimers.install({ toFake: ['setTimeout', 'clearTimeout'] })
+    t.teardown(() => clock.uninstall())
+
+    let count = 0
+    const MockConnection = connection.buildMockConnection({
+      onRequest (params) {
+        t.true(count < 2)
+        t.strictEqual(params.path, '/_bulk')
+        t.match(params.headers, { 'content-type': 'application/x-ndjson' })
+        const [action, payload] = params.body.split('\n')
+        t.deepEqual(JSON.parse(action), { index: { _index: 'test' } })
+        t.deepEqual(JSON.parse(payload), dataset[count++])
+        return { body: { errors: false, items: [{}] } }
+      }
+    })
+
+    const client = new Client({
+      node: 'http://localhost:9200',
+      Connection: MockConnection
+    })
+
+    const b = client.helpers.bulk({
+      datasource: (async function * generator () {
+        for (const chunk of dataset) {
+          await clock.nextAsync()
+          if (chunk.user === 'tyrion') {
+            // Needed otherwise in Node.js 10
+            // the second request will never be sent
+            await Promise.resolve()
+            b.abort()
+          }
+          yield chunk
+        }
+      })(),
+      flushBytes: 5000000,
+      concurrency: 1,
+      onDocument (doc) {
+        return {
+          index: { _index: 'test' }
+        }
+      },
+      onDrop (doc) {
+        t.fail('This should never be called')
+      }
+    })
+
+    const result = await b
+
+    t.type(result.time, 'number')
+    t.type(result.bytes, 'number')
+    t.match(result, {
+      total: 2,
+      successful: 2,
+      retry: 0,
+      failed: 0,
+      aborted: true
+    })
+  })
+
+  t.end()
+})

test/unit/helpers/msearch.test.js

@@ -7,6 +7,7 @@
 const { test } = require('tap')
 const { Client, errors } = require('../../../')
 const { connection } = require('../../utils')
+const FakeTimers = require('@sinonjs/fake-timers')
 
 test('Basic', async t => {
   const MockConnection = connection.buildMockConnection({
@@ -578,3 +579,158 @@ test('Multiple searches (concurrency = 1)', t => {
 
   t.teardown(() => s.stop())
 })
+
+test('Flush interval', t => {
+  t.plan(4)
+  const clock = FakeTimers.install({ toFake: ['setTimeout', 'clearTimeout'] })
+  t.teardown(() => clock.uninstall())
+
+  const MockConnection = connection.buildMockConnection({
+    onRequest (params) {
+      return {
+        body: {
+          responses: [{
+            status: 200,
+            hits: {
+              hits: [
+                { _source: { one: 'one' } },
+                { _source: { two: 'two' } },
+                { _source: { three: 'three' } }
+              ]
+            }
+          }, {
+            status: 200,
+            hits: {
+              hits: [
+                { _source: { four: 'four' } },
+                { _source: { five: 'five' } },
+                { _source: { six: 'six' } }
+              ]
+            }
+          }]
+        }
+      }
+    }
+  })
+
+  const client = new Client({
+    node: 'http://localhost:9200',
+    Connection: MockConnection
+  })
+
+  const s = client.helpers.msearch()
+
+  s.search({ index: 'test' }, { query: { match: { foo: 'bar' } } }, (err, result) => {
+    t.error(err)
+    t.is(result.documents.length, 3)
+  })
+
+  s.search({ index: 'test' }, { query: { match: { foo: 'bar' } } }, (err, result) => {
+    t.error(err)
+    t.is(result.documents.length, 3)
+  })
+
+  setImmediate(clock.next)
+
+  t.teardown(() => s.stop())
+})
+
+test('Flush interval - early stop', t => {
+  t.plan(3)
+  const clock = FakeTimers.install({ toFake: ['setTimeout', 'clearTimeout'] })
+  t.teardown(() => clock.uninstall())
+
+  const MockConnection = connection.buildMockConnection({
+    onRequest (params) {
+      return {
+        body: {
+          responses: [{
+            status: 200,
+            hits: {
+              hits: [
+                { _source: { one: 'one' } },
+                { _source: { two: 'two' } },
+                { _source: { three: 'three' } }
+              ]
+            }
+          }]
+        }
+      }
+    }
+  })
+
+  const client = new Client({
+    node: 'http://localhost:9200',
+    Connection: MockConnection
+  })
+
+  const s = client.helpers.msearch()
+
+  s.search({ index: 'test' }, { query: { match: { foo: 'bar' } } }, (err, result) => {
+    t.error(err)
+    t.is(result.documents.length, 3)
+  })
+
+  setImmediate(() => {
+    clock.next()
+    s.search({ index: 'test' }, { query: { match: { foo: 'bar' } } }, (err, result) => {
+      t.ok(err instanceof errors.ConfigurationError)
+    })
+  })
+
+  s.stop()
+})
+
+test('Stop should resolve the helper', t => {
+  t.plan(1)
+
+  const MockConnection = connection.buildMockConnection({
+    onRequest (params) {
+      return {
+        body: {
+          responses: []
+        }
+      }
+    }
+  })
+
+  const client = new Client({
+    node: 'http://localhost:9200',
+    Connection: MockConnection
+  })
+
+  const s = client.helpers.msearch()
+  setImmediate(s.stop)
+
+  s.then(() => t.pass('Called'))
+    .catch(() => t.fail('Should not fail'))
+})
+
+test('Stop should resolve the helper (error)', t => {
+  t.plan(3)
+
+  const MockConnection = connection.buildMockConnection({
+    onRequest (params) {
+      return {
+        body: {
+          responses: []
+        }
+      }
+    }
+  })
+
+  const client = new Client({
+    node: 'http://localhost:9200',
+    Connection: MockConnection
+  })
+
+  const s = client.helpers.msearch()
+  setImmediate(s.stop, new Error('kaboom'))
+
+  s.then(() => t.fail('Should not fail'))
+    .catch(err => t.is(err.message, 'kaboom'))
+
+  s.catch(err => t.is(err.message, 'kaboom'))
+
+  s.then(() => t.fail('Should not fail'), err => t.is(err.message, 'kaboom'))
+})