fix: ensure identify streams are closed

Author: jacobheun

src/identify/index.js

@@ -4,7 +4,7 @@ const debug = require('debug')
 const pb = require('it-protocol-buffers')
 const lp = require('it-length-prefixed')
 const pipe = require('it-pipe')
-const { collect, take } = require('streaming-iterables')
+const { collect, take, consume } = require('streaming-iterables')
 
 const PeerInfo = require('peer-info')
 const PeerId = require('peer-id')
@@ -114,7 +114,8 @@ class IdentifyService {
             protocols: Array.from(this._protocols.keys())
           }],
           pb.encode(Message),
-          stream
+          stream,
+          consume
         )
       } catch (err) {
         // Just log errors
@@ -159,6 +160,8 @@ class IdentifyService {
       toBuffer,
       collect
     )
+    // close the stream, no need to wait
+    stream.sink([])
 
     if (!data) {
       throw errCode(new Error('No data could be retrieved'), codes.ERR_CONNECTION_ENDED)
@@ -242,7 +245,8 @@ class IdentifyService {
     pipe(
       [message],
       lp.encode(),
-      stream
+      stream,
+      consume
     )
   }
 
@@ -261,6 +265,8 @@ class IdentifyService {
       toBuffer,
       collect
     )
+    // close the stream, but no need to wait
+    stream.sink([])
 
     let message
     try {

test/identify/index.spec.js

@@ -12,6 +12,7 @@ const PeerId = require('peer-id')
 const PeerInfo = require('peer-info')
 const duplexPair = require('it-pair/duplex')
 const multiaddr = require('multiaddr')
+const pWaitFor = require('p-wait-for')
 
 const { codes: Errors } = require('../../src/errors')
 const { IdentifyService, multicodecs } = require('../../src/identify')
@@ -203,16 +204,17 @@ describe('Identify', () => {
       })
 
       sinon.spy(libp2p.identifyService, 'identify')
-      sinon.spy(libp2p.peerStore, 'replace')
+      const peerStoreSpy = sinon.spy(libp2p.peerStore, 'replace')
 
       const connection = await libp2p.dialer.connectToPeer(remoteAddr)
       expect(connection).to.exist()
-      // Wait for nextTick to trigger the identify call
-      await delay(1)
+
+      // Wait for peer store to be updated
+      await pWaitFor(() => peerStoreSpy.callCount === 1)
       expect(libp2p.identifyService.identify.callCount).to.equal(1)
-      await libp2p.identifyService.identify.firstCall.returnValue
 
-      expect(libp2p.peerStore.replace.callCount).to.equal(1)
+      // The connection should have no open streams
+      expect(connection.streams).to.have.length(0)
       await connection.close()
     })
 
@@ -247,6 +249,9 @@ describe('Identify', () => {
         const results = await call.returnValue
         expect(results.length).to.equal(1)
       }
+
+      // The connection should have no open streams
+      expect(connection.streams).to.have.length(0)
     })
   })
 })