stream: Handle multi-corking properly

This adds proper support for the following situation:

    w.cork();
    w.write(...);
    w.cork();
    w.write(...);
    w.uncork();
    w.write(...);
    w.uncork();

This is relevant when you have a function (as we do in HTTP) that wants
to use cork, but in some cases, want to have a cork/uncork *around*
that function, without losing the benefits of writev.

1 files changed, 11 insertions, 5 deletions

diff --git a/lib/_stream_writable.js b/lib/_stream_writable.js
index f0034e3ba5..e40239ba86 100644
--- a/lib/_stream_writable.js
+++ b/lib/_stream_writable.js
@@ -77,7 +77,7 @@ function WritableState(options, stream) {
   this.writing = false;
 
   // when true all writes will be buffered until .uncork() call
-  this.corked = false;
+  this.corked = 0;
 
   // a flag to be able to tell if the onwrite cb is called immediately,
   // or on a later tick.  We set this to true at first, becuase any
@@ -190,16 +190,17 @@ Writable.prototype.write = function(chunk, encoding, cb) {
 Writable.prototype.cork = function() {
   var state = this._writableState;
 
-  state.corked = true;
+  state.corked++;
 };
 
 Writable.prototype.uncork = function() {
   var state = this._writableState;
 
   if (state.corked) {
-    state.corked = false;
+    state.corked--;
 
     if (!state.writing &&
+        !state.corked &&
         !state.finished &&
         !state.bufferProcessing &&
         state.buffer.length)
@@ -221,6 +222,8 @@ function decodeChunk(state, chunk, encoding) {
 // If we return false, then we need a drain event, so set that flag.
 function writeOrBuffer(stream, state, chunk, encoding, cb) {
   chunk = decodeChunk(state, chunk, encoding);
+  if (Buffer.isBuffer(chunk))
+    encoding = 'buffer';
   var len = state.objectMode ? 1 : chunk.length;
 
   state.length += len;
@@ -392,8 +395,11 @@ Writable.prototype.end = function(chunk, encoding, cb) {
   if (typeof chunk !== 'undefined' && chunk !== null)
     this.write(chunk, encoding);
 
-  // .end() should .uncork()
-  this.uncork();
+  // .end() fully uncorks
+  if (state.corked) {
+    state.corked = 1;
+    this.uncork();
+  }
 
   // ignore unnecessary end() calls.
   if (!state.ending && !state.finished)