feat: improved stream handling

__tests__/application/respond.test.js

@@ -943,6 +943,212 @@ describe('app.respond', () => {
 
       assert(streamDestroyed, 'Stream should be destroyed on client abort')
     })
+
+    it('should not emit error on premature close from client disconnect', async () => {
+      const app = new Koa()
+      const PassThrough = require('stream').PassThrough
+      const http = require('http')
+
+      let errorEmitted = false
+      app.on('error', () => {
+        errorEmitted = true
+      })
+
+      app.use(ctx => {
+        const stream = new PassThrough()
+        ctx.body = stream
+
+        setImmediate(() => {
+          stream.write('some data')
+        })
+      })
+
+      const server = app.listen()
+
+      await new Promise((resolve) => {
+        const req = http.request({
+          port: server.address().port,
+          path: '/'
+        })
+
+        req.on('response', (res) => {
+          res.on('data', () => {
+            req.destroy()
+            setTimeout(() => {
+              server.close()
+              resolve()
+            }, 50)
+          })
+        })
+
+        req.end()
+      })
+
+      assert.strictEqual(errorEmitted, false, 'ERR_STREAM_PREMATURE_CLOSE should not be emitted as an error')
+    })
+  })
+
+  describe('when .body is an AsyncIterable', () => {
+    it('should respond with async generator content', async () => {
+      const app = new Koa()
+
+      app.use(ctx => {
+        ctx.type = 'text/plain'
+        async function * generate () {
+          yield 'Hello '
+          yield 'World'
+        }
+        ctx.body = generate()
+      })
+
+      const res = await request(app.callback())
+        .get('/')
+        .expect(200)
+
+      assert.strictEqual(res.text, 'Hello World')
+    })
+
+    it('should respond with async iterable content', async () => {
+      const app = new Koa()
+
+      app.use(ctx => {
+        ctx.type = 'text/plain'
+        const iterable = {
+          [Symbol.asyncIterator] () {
+            let i = 0
+            const chunks = ['chunk1', 'chunk2', 'chunk3']
+            return {
+              async next () {
+                if (i < chunks.length) {
+                  return { value: chunks[i++], done: false }
+                }
+                return { done: true }
+              }
+            }
+          }
+        }
+        ctx.body = iterable
+      })
+
+      const res = await request(app.callback())
+        .get('/')
+        .expect(200)
+
+      assert.strictEqual(res.text, 'chunk1chunk2chunk3')
+    })
+
+    it('should respond with async generator yielding buffers', async () => {
+      const app = new Koa()
+
+      app.use(ctx => {
+        async function * generate () {
+          yield Buffer.from('Hello ')
+          yield Buffer.from('World')
+        }
+        ctx.body = generate()
+      })
+
+      const res = await request(app.callback())
+        .get('/')
+        .expect(200)
+
+      assert.deepStrictEqual(res.body, Buffer.from('Hello World'))
+    })
+
+    it('should default to octet-stream content type', async () => {
+      const app = new Koa()
+
+      app.use(ctx => {
+        async function * generate () {
+          yield 'data'
+        }
+        ctx.body = generate()
+      })
+
+      return request(app.callback())
+        .get('/')
+        .expect(200)
+        .expect('content-type', 'application/octet-stream')
+    })
+
+    it('should respect custom content type', async () => {
+      const app = new Koa()
+
+      app.use(ctx => {
+        ctx.type = 'text/plain'
+        async function * generate () {
+          yield 'plain text content'
+        }
+        ctx.body = generate()
+      })
+
+      return request(app.callback())
+        .get('/')
+        .expect(200)
+        .expect('content-type', 'text/plain; charset=utf-8')
+    })
+
+    it('should strip content-length when overwriting body with async iterable', async () => {
+      const app = new Koa()
+
+      app.use(ctx => {
+        ctx.body = 'hello'
+        async function * generate () {
+          yield 'async content'
+        }
+        ctx.body = generate()
+      })
+
+      const res = await request(app.callback())
+        .get('/')
+        .expect(200)
+
+      assert.strictEqual(Object.prototype.hasOwnProperty.call(res.headers, 'content-length'), false)
+      assert.strictEqual(res.text, 'async content')
+    })
+
+    it('should handle async generator errors', async () => {
+      const app = new Koa()
+
+      let errorCaught = false
+      app.once('error', err => {
+        assert.strictEqual(err.message, 'generator error')
+        errorCaught = true
+      })
+
+      app.use(ctx => {
+        async function * generate () {
+          yield 'start'
+          throw new Error('generator error')
+        }
+        ctx.body = generate()
+      })
+
+      await request(app.callback())
+        .get('/')
+        .catch(() => {})
+
+      await new Promise(resolve => setTimeout(resolve, 50))
+      assert(errorCaught, 'Error should have been caught')
+    })
+
+    it('should handle empty async generator', async () => {
+      const app = new Koa()
+
+      app.use(ctx => {
+        ctx.type = 'text/plain'
+        async function * generate () {
+          // empty
+        }
+        ctx.body = generate()
+      })
+
+      const res = await request(app.callback())
+        .get('/')
+        .expect(200)
+
+      assert.strictEqual(res.text, '')
+    })
   })
 
   describe('when .body is an Object', () => {

__tests__/response/body.test.js

@@ -304,6 +304,51 @@ describe('res.body=', () => {
     })
   })
 
+  describe('when an async iterable is given', () => {
+    it('should default to an octet stream', () => {
+      const res = response()
+      async function * generate () {
+        yield 'data'
+      }
+      res.body = generate()
+      assert.strictEqual('application/octet-stream', res.header['content-type'])
+    })
+
+    it('should not override Content-Type if already set', () => {
+      const res = response()
+      res.type = 'text/plain'
+      async function * generate () {
+        yield 'data'
+      }
+      res.body = generate()
+      assert.strictEqual('text/plain; charset=utf-8', res.header['content-type'])
+    })
+
+    it('should strip Content-Length when overwriting body', () => {
+      const res = response()
+      res.body = 'hello'
+      assert.strictEqual(5, res.header['content-length'])
+      async function * generate () {
+        yield 'data'
+      }
+      res.body = generate()
+      assert.strictEqual(res.header['content-length'], undefined)
+    })
+
+    it('should cleanup previous stream when replaced by async iterable', () => {
+      const res = response()
+      const stream = new Stream.PassThrough()
+
+      res.body = stream
+      async function * generate () {
+        yield 'data'
+      }
+      res.body = generate()
+
+      assert.strictEqual(stream.destroyed, true)
+    })
+  })
+
   describe('when a response is given', () => {
     it('should set the status', () => {
       const res = response()

lib/application.js

@@ -310,9 +310,13 @@ function respond (ctx) {
   else if (body instanceof ReadableStream) stream = Stream.Readable.from(body)
   else if (body instanceof Response) stream = Stream.Readable.from(body?.body || '')
   else if (isStream(body)) stream = body
+  else if (typeof body[Symbol.asyncIterator] === 'function') stream = Stream.Readable.from(body)
 
   if (stream) {
     return Stream.pipeline(stream, res, err => {
+      // Filter out premature close errors caused by client disconnects,
+      // as these are expected and not application errors.
+      if (err && err.code === 'ERR_STREAM_PREMATURE_CLOSE') return
       if (err && ctx.app.listenerCount('error')) ctx.onerror(err)
     })
   }

lib/is-stream.js

@@ -4,7 +4,7 @@ const Stream = require('stream')
 
 // TODO: use a third party library for this
 
-module.exports = (stream) => {
+const isStream = (stream) => {
   return (
     stream instanceof Stream ||
     (stream !== null &&
@@ -18,3 +18,5 @@ module.exports = (stream) => {
       typeof stream.destroyed === 'boolean')
   )
 }
+
+module.exports = isStream

lib/response.js

@@ -225,6 +225,14 @@ module.exports = {
       return
     }
 
+    // async iterable (e.g. async generators)
+    if (typeof val[Symbol.asyncIterator] === 'function') {
+      if (original != null) this.remove('Content-Length')
+      if (setType) this.type = 'bin'
+      cleanupPreviousStream()
+      return
+    }
+
     // json
     this.remove('Content-Length')
     if (!this.type || !/\bjson\b/i.test(this.type)) this.type = 'json'