From 3cb9234702384322cde54e93c7646bb000310f85 Mon Sep 17 00:00:00 2001
From: Robert Nagy <ronagy@icloud.com>
Date: Fri, 25 Sep 2020 18:50:49 +0200
Subject: [PATCH] fix: Client.stream writableNeedDrain

Fixes: https://github.com/nodejs/undici/issues/441
Refs: https://github.com/nodejs/node/pull/35348
Refs: https://github.com/nodejs/node/issues/35341
---
 README.md            | 2 +-
 lib/client-stream.js | 6 ++++++
 lib/core/client.js   | 4 +++-
 lib/core/request.js  | 2 +-
 4 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index f9a768b988e..347d2e1f4f4 100644
--- a/README.md
+++ b/README.md
@@ -473,7 +473,7 @@ The `handler` parameter is defined as follow:
   * `statusCode: Number`
   * `headers: Array|Null`
   * `socket: Duplex`
-* `onHeaders(statusCode, headers, resume): Void`, invoked when statusCode and headers have been received.
+* `onHeaders(statusCode, headers, resume): Boolean`, invoked when statusCode and headers have been received.
   May be invoked multiple times due to 1xx informational headers.
   * `statusCode: Number`
   * `headers: Array|Null`, an array of key-value pairs. Keys are not automatically lowercased.
diff --git a/lib/client-stream.js b/lib/client-stream.js
index d3188cfd956..a1a0428fdc8 100644
--- a/lib/client-stream.js
+++ b/lib/client-stream.js
@@ -109,6 +109,12 @@ class StreamHandler extends AsyncResource {
     })
 
     this.res = res
+
+    const needDrain = res.writableNeedDrain !== undefined
+      ? res.writableNeedDrain
+      : res._writableState && res._writableState.needDrain
+
+    return needDrain === true
   }
 
   onData (chunk) {
diff --git a/lib/core/client.js b/lib/core/client.js
index 9b742674b4d..9a04af55f73 100644
--- a/lib/core/client.js
+++ b/lib/core/client.js
@@ -535,7 +535,9 @@ class Parser extends HTTPParser {
     }
 
     try {
-      request.onHeaders(statusCode, headers, statusCode < 200 ? null : socket[kResume])
+      if (request.onHeaders(statusCode, headers, statusCode < 200 ? null : socket[kResume]) === false) {
+        socket[kPause]()
+      }
     } catch (err) {
       util.destroy(socket, err)
       return 1
diff --git a/lib/core/request.js b/lib/core/request.js
index 742897422b5..475d5dff037 100644
--- a/lib/core/request.js
+++ b/lib/core/request.js
@@ -117,7 +117,7 @@ class Request {
 
     clearRequestTimeout(this)
 
-    this[kHandler].onHeaders(statusCode, headers, resume)
+    return this[kHandler].onHeaders(statusCode, headers, resume)
   }
 
   onBody (chunk, offset, length) {