From 2570a1ba30b4735c4214541171b4b7ebb25cf6cf Mon Sep 17 00:00:00 2001
From: Jacob Heun <jacobheun@gmail.com>
Date: Tue, 10 Dec 2019 14:02:18 +0100
Subject: [PATCH] fix: release tokens as soon as they are available

---
 src/dialer/dial-request.js        |  6 ++--
 test/dialing/dial-request.spec.js | 49 +++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/src/dialer/dial-request.js b/src/dialer/dial-request.js
index 7063de1a90..f27749292b 100644
--- a/src/dialer/dial-request.js
+++ b/src/dialer/dial-request.js
@@ -47,11 +47,12 @@ class DialRequest {
     const tokenHolder = new FIFO()
     tokens.forEach(token => tokenHolder.push(token))
     const dialAbortControllers = this.addrs.map(() => new AbortController())
-    let completedDials = 0
+    let startedDials = 0
 
     try {
       return await pAny(this.addrs.map(async (addr, i) => {
         const token = await tokenHolder.shift() // get token
+        startedDials++
         let conn
         try {
           const signal = dialAbortControllers[i].signal
@@ -59,9 +60,8 @@ class DialRequest {
           // Remove the successful AbortController so it is not aborted
           dialAbortControllers.splice(i, 1)
         } finally {
-          completedDials++
           // If we have more dials to make, recycle the token, otherwise release it
-          if (completedDials < this.addrs.length) {
+          if (startedDials < this.addrs.length) {
             tokenHolder.push(token)
           } else {
             this.dialer.releaseToken(tokens.splice(tokens.indexOf(token), 1)[0])
diff --git a/test/dialing/dial-request.spec.js b/test/dialing/dial-request.spec.js
index 821546289f..1a85dad4ad 100644
--- a/test/dialing/dial-request.spec.js
+++ b/test/dialing/dial-request.spec.js
@@ -11,6 +11,7 @@ const { AbortError } = require('libp2p-interfaces/src/transport/errors')
 const AbortController = require('abort-controller')
 const AggregateError = require('aggregate-error')
 const pDefer = require('p-defer')
+const delay = require('delay')
 
 const { DialRequest } = require('../../src/dialer/dial-request')
 const createMockConnection = require('../utils/mockConnection')
@@ -50,6 +51,54 @@ describe('Dial Request', () => {
     expect(dialer.releaseToken).to.have.property('callCount', tokens.length)
   })
 
+  it('should release tokens when all addr dials have started', async () => {
+    const mockConnection = await createMockConnection()
+    const deferred = pDefer()
+    const actions = {
+      1: async () => {
+        await delay(0)
+        return Promise.reject(error)
+      },
+      2: async () => {
+        await delay(0)
+        return Promise.reject(error)
+      },
+      3: () => deferred.promise
+    }
+    const dialAction = (num) => actions[num]()
+    const tokens = ['a', 'b']
+    const controller = new AbortController()
+    const dialer = {
+      getTokens: () => [...tokens],
+      releaseToken: () => {}
+    }
+
+    const dialRequest = new DialRequest({
+      addrs: Object.keys(actions),
+      dialer,
+      dialAction
+    })
+
+    sinon.spy(actions, 1)
+    sinon.spy(actions, 2)
+    sinon.spy(actions, 3)
+    sinon.spy(dialer, 'releaseToken')
+    dialRequest.run({ signal: controller.signal })
+    // Let the first dials run
+    await delay(10)
+
+    // Only 1 dial should remain, so 1 token should have been released
+    expect(actions[1]).to.have.property('callCount', 1)
+    expect(actions[2]).to.have.property('callCount', 1)
+    expect(actions[3]).to.have.property('callCount', 1)
+    expect(dialer.releaseToken).to.have.property('callCount', 1)
+
+    // Finish the dial
+    deferred.resolve(mockConnection)
+    await delay(0)
+    expect(dialer.releaseToken).to.have.property('callCount', 2)
+  })
+
   it('should throw an AggregateError if all dials fail', async () => {
     const actions = {
       1: () => Promise.reject(error),