Introduce WebSocketDeflatePredictor.

Introduce WebSocketDeflatePredictor, an interface class which determines
whether a WebSocketDeflateStream should compress the given message or
not.

BUG=163882

Review URL: https://codereview.chromium.org/39193005

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@231827 0039d316-1c4b-4281-b951-d872f2087c98

8 files changed, 532 insertions, 58 deletions

diff --git a/net/websockets/README b/net/websockets/README
index 0670903..ddd6f5b 100644
--- a/net/websockets/README
+++ b/net/websockets/README
@@ -35,6 +35,10 @@ websocket_basic_stream_test.cc
 websocket_channel.cc
 websocket_channel.h
 websocket_channel_test.cc
+websocket_deflate_predictor.h
+websocket_deflate_predictor_impl.cc
+websocket_deflate_predictor_impl.h
+websocket_deflate_predictor_impl_test.cc
 websocket_deflate_stream.cc
 websocket_deflate_stream.h
 websocket_deflate_stream_test.cc
diff --git a/net/websockets/websocket_deflate_predictor.h b/net/websockets/websocket_deflate_predictor.h
new file mode 100644
index 0000000..c786d80
--- /dev/null
+++ b/net/websockets/websocket_deflate_predictor.h
@@ -0,0 +1,58 @@
+// Copyright 2013 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef NET_WEBSOCKETS_WEBSOCKET_DEFLATE_PREDICTOR_H_
+#define NET_WEBSOCKETS_WEBSOCKET_DEFLATE_PREDICTOR_H_
+
+#include "base/basictypes.h"
+#include "base/memory/scoped_vector.h"
+#include "net/base/net_export.h"
+
+namespace net {
+
+struct WebSocketFrame;
+
+// WebSocketDeflatePredictor is an interface class used for judging whether
+// a WebSocketDeflateStream should compress a message or not.
+class NET_EXPORT_PRIVATE WebSocketDeflatePredictor {
+ public:
+  enum Result {
+    // Deflate and send the message.
+    DEFLATE,
+    // Do not deflate and send the original message.
+    DO_NOT_DEFLATE,
+    // Try compressing the message and send the smaller one of the original
+    // and the compressed message.
+    // Returning this result implies that the deflater is running on
+    // DoNotTakeOverContext mode and the entire message is visible.
+    TRY_DEFLATE,
+  };
+
+  virtual ~WebSocketDeflatePredictor() {}
+
+  // Predicts and returns whether the deflater should deflate the message
+  // which begins with |frames[frame_index]| or not.
+  // |frames[(frame_index + 1):]| consists of future frames if any.
+  // |frames[frame_index]| must be the first frame of a data message,
+  // but future frames may contain control message frames.
+  // |frames[frame_index]| cannot be recorded yet and all preceding
+  // data frames have to be already recorded when this method is called.
+  virtual Result Predict(const ScopedVector<WebSocketFrame>& frames,
+                         size_t frame_index) = 0;
+
+  // Records frame data for future prediction.
+  // Only data frames should be recorded. Do not pass control frames' data.
+  // All input data frames for the stream must be recorded in order.
+  virtual void RecordInputDataFrame(const WebSocketFrame* frame) = 0;
+
+  // Records frame data for future prediction.
+  // Only data frames should be recorded. Do not pass control frames' data.
+  // All data frames written by the stream must be recorded in order
+  // regardless of whether they are compressed or not.
+  virtual void RecordWrittenDataFrame(const WebSocketFrame* frame) = 0;
+};
+
+}  // namespace net
+
+#endif  // NET_WEBSOCKETS_WEBSOCKET_DEFLATE_PREDICTOR_H_
diff --git a/net/websockets/websocket_deflate_predictor_impl.cc b/net/websockets/websocket_deflate_predictor_impl.cc
new file mode 100644
index 0000000..0d1a5c20
--- /dev/null
+++ b/net/websockets/websocket_deflate_predictor_impl.cc
@@ -0,0 +1,23 @@
+// Copyright 2013 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "net/websockets/websocket_deflate_predictor_impl.h"
+
+namespace net {
+
+typedef WebSocketDeflatePredictor::Result Result;
+
+Result WebSocketDeflatePredictorImpl::Predict(
+    const ScopedVector<WebSocketFrame>& frames,
+    size_t frame_index) {
+  return DEFLATE;
+}
+
+void WebSocketDeflatePredictorImpl::RecordInputDataFrame(
+    const WebSocketFrame* frame) {}
+
+void WebSocketDeflatePredictorImpl::RecordWrittenDataFrame(
+    const WebSocketFrame* frame) {}
+
+}  // namespace net
diff --git a/net/websockets/websocket_deflate_predictor_impl.h b/net/websockets/websocket_deflate_predictor_impl.h
new file mode 100644
index 0000000..88a919c
--- /dev/null
+++ b/net/websockets/websocket_deflate_predictor_impl.h
@@ -0,0 +1,31 @@
+// Copyright 2013 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef NET_WEBSOCKETS_WEBSOCKET_DEFLATE_PREDICTOR_IMPL_H_
+#define NET_WEBSOCKETS_WEBSOCKET_DEFLATE_PREDICTOR_IMPL_H_
+
+#include "base/basictypes.h"
+#include "base/compiler_specific.h"
+#include "base/memory/scoped_vector.h"
+#include "net/base/net_export.h"
+#include "net/websockets/websocket_deflate_predictor.h"
+
+namespace net {
+
+struct WebSocketFrame;
+
+class NET_EXPORT_PRIVATE WebSocketDeflatePredictorImpl
+    : public WebSocketDeflatePredictor {
+ public:
+  virtual ~WebSocketDeflatePredictorImpl() {}
+
+  virtual Result Predict(const ScopedVector<WebSocketFrame>& frames,
+                         size_t frame_index) OVERRIDE;
+  virtual void RecordInputDataFrame(const WebSocketFrame* frame) OVERRIDE;
+  virtual void RecordWrittenDataFrame(const WebSocketFrame* frame) OVERRIDE;
+};
+
+}  // namespace net
+
+#endif  // NET_WEBSOCKETS_WEBSOCKET_DEFLATE_PREDICTOR_IMPL_H_
diff --git a/net/websockets/websocket_deflate_predictor_impl_test.cc b/net/websockets/websocket_deflate_predictor_impl_test.cc
new file mode 100644
index 0000000..79c54e1
--- /dev/null
+++ b/net/websockets/websocket_deflate_predictor_impl_test.cc
@@ -0,0 +1,29 @@
+// Copyright 2013 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "net/websockets/websocket_deflate_predictor_impl.h"
+
+#include "base/basictypes.h"
+#include "base/memory/scoped_vector.h"
+#include "net/websockets/websocket_frame.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+namespace net {
+
+namespace {
+
+typedef WebSocketDeflatePredictor::Result Result;
+
+TEST(WebSocketDeflatePredictorImpl, Predict) {
+  WebSocketDeflatePredictorImpl predictor;
+  ScopedVector<WebSocketFrame> frames;
+  frames.push_back(new WebSocketFrame(WebSocketFrameHeader::kOpCodeText));
+  Result result = predictor.Predict(frames, 0);
+
+  EXPECT_EQ(WebSocketDeflatePredictor::DEFLATE, result);
+}
+
+}  // namespace
+
+}  // namespace net
diff --git a/net/websockets/websocket_deflate_stream.cc b/net/websockets/websocket_deflate_stream.cc
index 9e8a95a..601670d 100644
--- a/net/websockets/websocket_deflate_stream.cc
+++ b/net/websockets/websocket_deflate_stream.cc
@@ -15,6 +15,7 @@
 #include "net/base/completion_callback.h"
 #include "net/base/io_buffer.h"
 #include "net/base/net_errors.h"
+#include "net/websockets/websocket_deflate_predictor.h"
 #include "net/websockets/websocket_deflater.h"
 #include "net/websockets/websocket_errors.h"
 #include "net/websockets/websocket_frame.h"
@@ -34,14 +35,16 @@ const size_t kChunkSize = 4 * 1024;
 
 WebSocketDeflateStream::WebSocketDeflateStream(
     scoped_ptr<WebSocketStream> stream,
-    WebSocketDeflater::ContextTakeOverMode mode)
+    WebSocketDeflater::ContextTakeOverMode mode,
+    scoped_ptr<WebSocketDeflatePredictor> predictor)
     : stream_(stream.Pass()),
       deflater_(mode),
       inflater_(kChunkSize, kChunkSize),
       reading_state_(NOT_READING),
       writing_state_(NOT_WRITING),
       current_reading_opcode_(WebSocketFrameHeader::kOpCodeText),
-      current_writing_opcode_(WebSocketFrameHeader::kOpCodeText) {
+      current_writing_opcode_(WebSocketFrameHeader::kOpCodeText),
+      predictor_(predictor.Pass()) {
   DCHECK(stream_);
   deflater_.Initialize(kWindowBits);
   inflater_.Initialize(kWindowBits);
@@ -100,31 +103,30 @@ void WebSocketDeflateStream::OnReadComplete(
 
 int WebSocketDeflateStream::Deflate(ScopedVector<WebSocketFrame>* frames) {
   ScopedVector<WebSocketFrame> frames_to_write;
+  // Store frames of the currently processed message if writing_state_ equals to
+  // WRITING_POSSIBLY_COMPRESSED_MESSAGE.
+  ScopedVector<WebSocketFrame> frames_of_message;
   for (size_t i = 0; i < frames->size(); ++i) {
-    scoped_ptr<WebSocketFrame> frame((*frames)[i]);
-    (*frames)[i] = NULL;
-    DCHECK(!frame->header.reserved1);
-    if (!WebSocketFrameHeader::IsKnownDataOpCode(frame->header.opcode)) {
-      frames_to_write.push_back(frame.release());
+    DCHECK(!(*frames)[i]->header.reserved1);
+    if (!WebSocketFrameHeader::IsKnownDataOpCode((*frames)[i]->header.opcode)) {
+      frames_to_write.push_back((*frames)[i]);
+      (*frames)[i] = NULL;
       continue;
     }
+    if (writing_state_ == NOT_WRITING)
+      OnMessageStart(*frames, i);
+
+    scoped_ptr<WebSocketFrame> frame((*frames)[i]);
+    (*frames)[i] = NULL;
+    predictor_->RecordInputDataFrame(frame.get());
 
-    if (writing_state_ == NOT_WRITING) {
-      current_writing_opcode_ = frame->header.opcode;
-      DCHECK(current_writing_opcode_ == WebSocketFrameHeader::kOpCodeText ||
-             current_writing_opcode_ == WebSocketFrameHeader::kOpCodeBinary);
-      // TODO(yhirano): For now, we unconditionally compress data messages.
-      // Further optimization is needed.
-      // http://crbug.com/163882
-      writing_state_ = WRITING_COMPRESSED_MESSAGE;
-    }
     if (writing_state_ == WRITING_UNCOMPRESSED_MESSAGE) {
       if (frame->header.final)
         writing_state_ = NOT_WRITING;
+      predictor_->RecordWrittenDataFrame(frame.get());
       frames_to_write.push_back(frame.release());
       current_writing_opcode_ = WebSocketFrameHeader::kOpCodeContinuation;
     } else {
-      DCHECK_EQ(WRITING_COMPRESSED_MESSAGE, writing_state_);
       if (frame->data && !deflater_.AddBytes(frame->data->data(),
                                              frame->header.payload_length)) {
         DVLOG(1) << "WebSocket protocol error. "
@@ -136,35 +138,135 @@ int WebSocketDeflateStream::Deflate(ScopedVector<WebSocketFrame>* frames) {
                  << "deflater_.Finish() returns an error.";
         return ERR_WS_PROTOCOL_ERROR;
       }
-      if (deflater_.CurrentOutputSize() >= kChunkSize || frame->header.final) {
-        const WebSocketFrameHeader::OpCode opcode = current_writing_opcode_;
-        scoped_ptr<WebSocketFrame> compressed(new WebSocketFrame(opcode));
-        scoped_refptr<IOBufferWithSize> data =
-            deflater_.GetOutput(deflater_.CurrentOutputSize());
-        if (!data) {
-          DVLOG(1) << "WebSocket protocol error. "
-                   << "deflater_.GetOutput() returns an error.";
-          return ERR_WS_PROTOCOL_ERROR;
+
+      if (writing_state_ == WRITING_COMPRESSED_MESSAGE) {
+        if (deflater_.CurrentOutputSize() >= kChunkSize ||
+            frame->header.final) {
+          int result = AppendCompressedFrame(frame->header, &frames_to_write);
+          if (result != OK)
+            return result;
+        }
+        if (frame->header.final)
+          writing_state_ = NOT_WRITING;
+      } else {
+        DCHECK_EQ(WRITING_POSSIBLY_COMPRESSED_MESSAGE, writing_state_);
+        bool final = frame->header.final;
+        frames_of_message.push_back(frame.release());
+        if (final) {
+          int result = AppendPossiblyCompressedMessage(&frames_of_message,
+                                                       &frames_to_write);
+          if (result != OK)
+            return result;
+          frames_of_message.clear();
+          writing_state_ = NOT_WRITING;
         }
-        compressed->header.CopyFrom(frame->header);
-        compressed->header.opcode = opcode;
-        compressed->header.final = frame->header.final;
-        compressed->header.reserved1 =
-            (opcode != WebSocketFrameHeader::kOpCodeContinuation);
-        compressed->data = data;
-        compressed->header.payload_length = data->size();
-
-        current_writing_opcode_ = WebSocketFrameHeader::kOpCodeContinuation;
-        frames_to_write.push_back(compressed.release());
       }
-      if (frame->header.final)
-        writing_state_ = NOT_WRITING;
     }
   }
+  DCHECK_NE(WRITING_POSSIBLY_COMPRESSED_MESSAGE, writing_state_);
   frames->swap(frames_to_write);
   return OK;
 }
 
+void WebSocketDeflateStream::OnMessageStart(
+    const ScopedVector<WebSocketFrame>& frames, size_t index) {
+  WebSocketFrame* frame = frames[index];
+  current_writing_opcode_ = frame->header.opcode;
+  DCHECK(current_writing_opcode_ == WebSocketFrameHeader::kOpCodeText ||
+         current_writing_opcode_ == WebSocketFrameHeader::kOpCodeBinary);
+  WebSocketDeflatePredictor::Result prediction =
+      predictor_->Predict(frames, index);
+
+  switch (prediction) {
+    case WebSocketDeflatePredictor::DEFLATE:
+      writing_state_ = WRITING_COMPRESSED_MESSAGE;
+      return;
+    case WebSocketDeflatePredictor::DO_NOT_DEFLATE:
+      writing_state_ = WRITING_UNCOMPRESSED_MESSAGE;
+      return;
+    case WebSocketDeflatePredictor::TRY_DEFLATE:
+      writing_state_ = WRITING_POSSIBLY_COMPRESSED_MESSAGE;
+      return;
+  }
+  NOTREACHED();
+}
+
+int WebSocketDeflateStream::AppendCompressedFrame(
+    const WebSocketFrameHeader& header,
+    ScopedVector<WebSocketFrame>* frames_to_write) {
+  const WebSocketFrameHeader::OpCode opcode = current_writing_opcode_;
+  scoped_refptr<IOBufferWithSize> compressed_payload =
+      deflater_.GetOutput(deflater_.CurrentOutputSize());
+  if (!compressed_payload) {
+    DVLOG(1) << "WebSocket protocol error. "
+             << "deflater_.GetOutput() returns an error.";
+    return ERR_WS_PROTOCOL_ERROR;
+  }
+  scoped_ptr<WebSocketFrame> compressed(new WebSocketFrame(opcode));
+  compressed->header.CopyFrom(header);
+  compressed->header.opcode = opcode;
+  compressed->header.final = header.final;
+  compressed->header.reserved1 =
+      (opcode != WebSocketFrameHeader::kOpCodeContinuation);
+  compressed->data = compressed_payload;
+  compressed->header.payload_length = compressed_payload->size();
+
+  current_writing_opcode_ = WebSocketFrameHeader::kOpCodeContinuation;
+  predictor_->RecordWrittenDataFrame(compressed.get());
+  frames_to_write->push_back(compressed.release());
+  return OK;
+}
+
+int WebSocketDeflateStream::AppendPossiblyCompressedMessage(
+    ScopedVector<WebSocketFrame>* frames,
+    ScopedVector<WebSocketFrame>* frames_to_write) {
+  DCHECK(!frames->empty());
+
+  const WebSocketFrameHeader::OpCode opcode = current_writing_opcode_;
+  scoped_refptr<IOBufferWithSize> compressed_payload =
+      deflater_.GetOutput(deflater_.CurrentOutputSize());
+  if (!compressed_payload) {
+    DVLOG(1) << "WebSocket protocol error. "
+             << "deflater_.GetOutput() returns an error.";
+    return ERR_WS_PROTOCOL_ERROR;
+  }
+
+  uint64 original_payload_length = 0;
+  for (size_t i = 0; i < frames->size(); ++i) {
+    WebSocketFrame* frame = (*frames)[i];
+    // Asserts checking that frames represent one whole data message.
+    DCHECK(WebSocketFrameHeader::IsKnownDataOpCode(frame->header.opcode));
+    DCHECK_EQ(i == 0,
+              WebSocketFrameHeader::kOpCodeContinuation !=
+              frame->header.opcode);
+    DCHECK_EQ(i == frames->size() - 1, frame->header.final);
+    original_payload_length += frame->header.payload_length;
+  }
+  if (original_payload_length <=
+      static_cast<uint64>(compressed_payload->size())) {
+    // Compression is not effective. Use the original frames.
+    for (size_t i = 0; i < frames->size(); ++i) {
+      WebSocketFrame* frame = (*frames)[i];
+      frames_to_write->push_back(frame);
+      predictor_->RecordWrittenDataFrame(frame);
+      (*frames)[i] = NULL;
+    }
+    frames->weak_clear();
+    return OK;
+  }
+  scoped_ptr<WebSocketFrame> compressed(new WebSocketFrame(opcode));
+  compressed->header.CopyFrom((*frames)[0]->header);
+  compressed->header.opcode = opcode;
+  compressed->header.final = true;
+  compressed->header.reserved1 = true;
+  compressed->data = compressed_payload;
+  compressed->header.payload_length = compressed_payload->size();
+
+  predictor_->RecordWrittenDataFrame(compressed.get());
+  frames_to_write->push_back(compressed.release());
+  return OK;
+}
+
 int WebSocketDeflateStream::Inflate(ScopedVector<WebSocketFrame>* frames) {
   ScopedVector<WebSocketFrame> frames_to_output;
   ScopedVector<WebSocketFrame> frames_passed;
diff --git a/net/websockets/websocket_deflate_stream.h b/net/websockets/websocket_deflate_stream.h
index fc2ebf4..a785944 100644
--- a/net/websockets/websocket_deflate_stream.h
+++ b/net/websockets/websocket_deflate_stream.h
@@ -21,14 +21,27 @@ class GURL;
 
 namespace net {
 
+class WebSocketDeflatePredictor;
+
 // WebSocketDeflateStream is a WebSocketStream subclass.
 // WebSocketDeflateStream is for permessage-deflate WebSocket extension[1].
 //
+// WebSocketDeflateStream::ReadFrames and WriteFrames may change frame
+// boundary. In particular, if a control frame is placed in the middle of
+// data message frames, the control frame can overtake data frames.
+// Say there are frames df1, df2 and cf, df1 and df2 are frames of a
+// data message and cf is a control message frame. cf may arrive first and
+// data frames may follow cf.
+// Note that message boundary will be preserved, i.e. if the last frame of
+// a message m1 is read / written before the last frame of a message m2,
+// WebSocketDeflateStream will respect the order.
+//
 // [1]: http://tools.ietf.org/html/draft-ietf-hybi-permessage-compression-12
 class NET_EXPORT_PRIVATE WebSocketDeflateStream : public WebSocketStream {
  public:
   WebSocketDeflateStream(scoped_ptr<WebSocketStream> stream,
-                         WebSocketDeflater::ContextTakeOverMode mode);
+                         WebSocketDeflater::ContextTakeOverMode mode,
+                         scoped_ptr<WebSocketDeflatePredictor> predictor);
   virtual ~WebSocketDeflateStream();
 
   // WebSocketStream functions.
@@ -50,6 +63,7 @@ class NET_EXPORT_PRIVATE WebSocketDeflateStream : public WebSocketStream {
   enum WritingState {
     WRITING_COMPRESSED_MESSAGE,
     WRITING_UNCOMPRESSED_MESSAGE,
+    WRITING_POSSIBLY_COMPRESSED_MESSAGE,
     NOT_WRITING,
   };
 
@@ -59,6 +73,12 @@ class NET_EXPORT_PRIVATE WebSocketDeflateStream : public WebSocketStream {
 
   // This function deflates |frames| and stores the result to |frames| itself.
   int Deflate(ScopedVector<WebSocketFrame>* frames);
+  void OnMessageStart(const ScopedVector<WebSocketFrame>& frames, size_t index);
+  int AppendCompressedFrame(const WebSocketFrameHeader& header,
+                            ScopedVector<WebSocketFrame>* frames_to_write);
+  int AppendPossiblyCompressedMessage(
+      ScopedVector<WebSocketFrame>* frames,
+      ScopedVector<WebSocketFrame>* frames_to_write);
 
   // This function inflates |frames| and stores the result to |frames| itself.
   int Inflate(ScopedVector<WebSocketFrame>* frames);
@@ -73,6 +93,7 @@ class NET_EXPORT_PRIVATE WebSocketDeflateStream : public WebSocketStream {
   WritingState writing_state_;
   WebSocketFrameHeader::OpCode current_reading_opcode_;
   WebSocketFrameHeader::OpCode current_writing_opcode_;
+  scoped_ptr<WebSocketDeflatePredictor> predictor_;
 
   DISALLOW_COPY_AND_ASSIGN(WebSocketDeflateStream);
 };
diff --git a/net/websockets/websocket_deflate_stream_test.cc b/net/websockets/websocket_deflate_stream_test.cc
index 09b2c5c..1775962 100644
--- a/net/websockets/websocket_deflate_stream_test.cc
+++ b/net/websockets/websocket_deflate_stream_test.cc
@@ -5,6 +5,7 @@
 #include "net/websockets/websocket_deflate_stream.h"
 
 #include <stdint.h>
+#include <deque>
 #include <string>
 
 #include "base/basictypes.h"
@@ -15,6 +16,7 @@
 #include "net/base/completion_callback.h"
 #include "net/base/io_buffer.h"
 #include "net/base/net_errors.h"
+#include "net/websockets/websocket_deflate_predictor.h"
 #include "net/websockets/websocket_deflater.h"
 #include "net/websockets/websocket_frame.h"
 #include "net/websockets/websocket_inflater.h"
@@ -65,7 +67,7 @@ std::string ToString(const scoped_refptr<IOBuffer>& buffer, size_t size) {
   return ToString(buffer.get(), size);
 }
 
-std::string ToString(WebSocketFrame* frame) {
+std::string ToString(const WebSocketFrame* frame) {
   return frame->data ? ToString(frame->data, frame->header.payload_length) : "";
 }
 
@@ -101,25 +103,132 @@ class MockWebSocketStream : public WebSocketStream {
   MOCK_CONST_METHOD0(GetExtensions, std::string());
 };
 
+// This mock class relies on some assumptions.
+//  - RecordInputDataFrame is called after the corresponding WriteFrames
+//    call.
+//  - RecordWrittenDataFrame is called before writing the frame.
+class WebSocketDeflatePredictorMock : public WebSocketDeflatePredictor {
+ public:
+  WebSocketDeflatePredictorMock() : result_(DEFLATE) {}
+  virtual ~WebSocketDeflatePredictorMock() {
+    // Verify whether all expectaions are consumed.
+    if (!frames_to_be_input_.empty()) {
+      ADD_FAILURE() << "There are missing frames to be input.";
+      return;
+    }
+    if (!frames_written_.empty()) {
+      ADD_FAILURE() << "There are extra written frames.";
+      return;
+    }
+  }
+
+  // WebSocketDeflatePredictor functions.
+  virtual Result Predict(const ScopedVector<WebSocketFrame>& frames,
+                         size_t frame_index) OVERRIDE {
+    return result_;
+  }
+  virtual void RecordInputDataFrame(const WebSocketFrame* frame) OVERRIDE {
+    if (!WebSocketFrameHeader::IsKnownDataOpCode(frame->header.opcode)) {
+      ADD_FAILURE() << "Control frames should not be recorded.";
+      return;
+    }
+    if (frame->header.reserved1) {
+      ADD_FAILURE() << "Input frame may not be compressed.";
+      return;
+    }
+    if (frames_to_be_input_.empty()) {
+      ADD_FAILURE() << "Unexpected input data frame";
+      return;
+    }
+    if (frame != frames_to_be_input_.front()) {
+      ADD_FAILURE() << "Input data frame does not match the expectation.";
+      return;
+    }
+    frames_to_be_input_.pop_front();
+  }
+  virtual void RecordWrittenDataFrame(const WebSocketFrame* frame) OVERRIDE {
+    if (!WebSocketFrameHeader::IsKnownDataOpCode(frame->header.opcode)) {
+      ADD_FAILURE() << "Control frames should not be recorded.";
+      return;
+    }
+    frames_written_.push_back(frame);
+  }
+
+  // Sets |result_| for the |Predict| return value.
+  void set_result(Result result) { result_ = result; }
+
+  // Adds |frame| as an expectation of future |RecordInputDataFrame| call.
+  void AddFrameToBeInput(const WebSocketFrame* frame) {
+    if (!WebSocketFrameHeader::IsKnownDataOpCode(frame->header.opcode))
+      return;
+    frames_to_be_input_.push_back(frame);
+  }
+  // Verifies that |frame| is recorded in order.
+  void VerifySentFrame(const WebSocketFrame* frame) {
+    if (!WebSocketFrameHeader::IsKnownDataOpCode(frame->header.opcode))
+      return;
+    if (frames_written_.empty()) {
+      ADD_FAILURE() << "There are missing frames to be written.";
+      return;
+    }
+    if (frame != frames_written_.front()) {
+      ADD_FAILURE() << "Written data frame does not match the expectation.";
+      return;
+    }
+    frames_written_.pop_front();
+  }
+  void AddFramesToBeInput(const ScopedVector<WebSocketFrame>& frames) {
+    for (size_t i = 0; i < frames.size(); ++i)
+      AddFrameToBeInput(frames[i]);
+  }
+  void VerifySentFrames(const ScopedVector<WebSocketFrame>& frames) {
+    for (size_t i = 0; i < frames.size(); ++i)
+      VerifySentFrame(frames[i]);
+  }
+  // Call this method in order to disable checks in the destructor when
+  // WriteFrames fails.
+  void Clear() {
+    frames_to_be_input_.clear();
+    frames_written_.clear();
+  }
+
+ private:
+  Result result_;
+  // Data frames which will be recorded by |RecordInputFrames|.
+  // Pushed by |AddFrameToBeInput| and popped and verified by
+  // |RecordInputFrames|.
+  std::deque<const WebSocketFrame*> frames_to_be_input_;
+  // Data frames recorded by |RecordWrittenFrames|.
+  // Pushed by |RecordWrittenFrames| and popped and verified by
+  // |VerifySentFrame|.
+  std::deque<const WebSocketFrame*> frames_written_;
+
+  DISALLOW_COPY_AND_ASSIGN(WebSocketDeflatePredictorMock);
+};
+
 class WebSocketDeflateStreamTest : public ::testing::Test {
  public:
   WebSocketDeflateStreamTest()
       : mock_stream_(NULL) {
     mock_stream_ = new testing::StrictMock<MockWebSocketStream>;
+    predictor_ = new WebSocketDeflatePredictorMock;
     deflate_stream_.reset(new WebSocketDeflateStream(
         scoped_ptr<WebSocketStream>(mock_stream_),
-        WebSocketDeflater::TAKE_OVER_CONTEXT));
+        WebSocketDeflater::TAKE_OVER_CONTEXT,
+        scoped_ptr<WebSocketDeflatePredictor>(predictor_)));
   }
   virtual ~WebSocketDeflateStreamTest() {}
 
  protected:
   scoped_ptr<WebSocketDeflateStream> deflate_stream_;
-  // |mock_stream_| will be deleted when |deflate_stream_| is destroyed.
+  // Owned by |deflate_stream_|.
   MockWebSocketStream* mock_stream_;
+  // Owned by |deflate_stream_|.
+  WebSocketDeflatePredictorMock* predictor_;
 };
 
 // Since WebSocketDeflater with DoNotTakeOverContext is well tested at
-// websocket_deflater_test.cc, we have only one test for this configuration
+// websocket_deflater_test.cc, we have only a few tests for this configuration
 // here.
 class WebSocketDeflateStreamWithDoNotTakeOverContextTest
     : public ::testing::Test {
@@ -127,9 +236,11 @@ class WebSocketDeflateStreamWithDoNotTakeOverContextTest
   WebSocketDeflateStreamWithDoNotTakeOverContextTest()
       : mock_stream_(NULL) {
     mock_stream_ = new testing::StrictMock<MockWebSocketStream>;
+    predictor_ = new WebSocketDeflatePredictorMock;
     deflate_stream_.reset(new WebSocketDeflateStream(
         scoped_ptr<WebSocketStream>(mock_stream_),
-        WebSocketDeflater::DO_NOT_TAKE_OVER_CONTEXT));
+        WebSocketDeflater::DO_NOT_TAKE_OVER_CONTEXT,
+        scoped_ptr<WebSocketDeflatePredictor>(predictor_)));
   }
   virtual ~WebSocketDeflateStreamWithDoNotTakeOverContextTest() {}
 
@@ -137,6 +248,8 @@ class WebSocketDeflateStreamWithDoNotTakeOverContextTest
   scoped_ptr<WebSocketDeflateStream> deflate_stream_;
   // |mock_stream_| will be deleted when |deflate_stream_| is destroyed.
   MockWebSocketStream* mock_stream_;
+  // |predictor_| will be deleted when |deflate_stream_| is destroyed.
+  WebSocketDeflatePredictorMock* predictor_;
 };
 
 // ReadFrameStub is a stub for WebSocketStream::ReadFrames.
@@ -173,20 +286,21 @@ class ReadFramesStub {
   ScopedVector<WebSocketFrame>* frames_passed_;
 };
 
-// WriteFrameStub is a stub for WebSocketStream::WriteFrames.
+// WriteFramesStub is a stub for WebSocketStream::WriteFrames.
 // It returns |result_| and |frames_| to the caller and
 // saves |callback| parameter to |callback_|.
 class WriteFramesStub {
  public:
-  explicit WriteFramesStub(int result) : result_(result) {}
+  explicit WriteFramesStub(WebSocketDeflatePredictorMock* predictor,
+                           int result)
+      : result_(result), predictor_(predictor) {}
 
   int Call(ScopedVector<WebSocketFrame>* frames,
            const CompletionCallback& callback) {
-    for (size_t i = 0; i < frames->size(); ++i) {
-      frames_.push_back((*frames)[i]);
-    }
+    frames_.insert(frames_.end(), frames->begin(), frames->end());
     frames->weak_clear();
     callback_ = callback;
+    predictor_->VerifySentFrames(frames_);
     return result_;
   }
 
@@ -198,6 +312,7 @@ class WriteFramesStub {
   int result_;
   CompletionCallback callback_;
   ScopedVector<WebSocketFrame> frames_;
+  WebSocketDeflatePredictorMock* predictor_;
 };
 
 TEST_F(WebSocketDeflateStreamTest, ReadFailedImmediately) {
@@ -791,14 +906,17 @@ TEST_F(WebSocketDeflateStreamTest, WriteFailedImmediately) {
   }
 
   AppendTo(&frames, WebSocketFrameHeader::kOpCodeText, kFinal, "hello");
+  predictor_->AddFramesToBeInput(frames);
   EXPECT_EQ(ERR_FAILED, deflate_stream_->WriteFrames(&frames, callback));
+  predictor_->Clear();
 }
 
 TEST_F(WebSocketDeflateStreamTest, WriteFrameImmediately) {
   ScopedVector<WebSocketFrame> frames;
   CompletionCallback callback;
-  WriteFramesStub stub(OK);
+  WriteFramesStub stub(predictor_, OK);
   AppendTo(&frames, WebSocketFrameHeader::kOpCodeText, kFinal, "Hello");
+  predictor_->AddFramesToBeInput(frames);
   {
     InSequence s;
     EXPECT_CALL(*mock_stream_, WriteFrames(_, _))
@@ -815,7 +933,7 @@ TEST_F(WebSocketDeflateStreamTest, WriteFrameImmediately) {
 }
 
 TEST_F(WebSocketDeflateStreamTest, WriteFrameAsync) {
-  WriteFramesStub stub(ERR_IO_PENDING);
+  WriteFramesStub stub(predictor_, ERR_IO_PENDING);
   MockCallback mock_callback, checkpoint;
   CompletionCallback callback =
       base::Bind(&MockCallback::Call, base::Unretained(&mock_callback));
@@ -828,6 +946,7 @@ TEST_F(WebSocketDeflateStreamTest, WriteFrameAsync) {
     EXPECT_CALL(mock_callback, Call(OK));
   }
   AppendTo(&frames, WebSocketFrameHeader::kOpCodeText, kFinal, "Hello");
+  predictor_->AddFramesToBeInput(frames);
   ASSERT_EQ(ERR_IO_PENDING, deflate_stream_->WriteFrames(&frames, callback));
 
   checkpoint.Call(0);
@@ -847,7 +966,8 @@ TEST_F(WebSocketDeflateStreamTest, WriteControlFrameBetweenDataFrames) {
   AppendTo(&frames, WebSocketFrameHeader::kOpCodeText, kNoFlag, "Hel");
   AppendTo(&frames, WebSocketFrameHeader::kOpCodePing, kFinal);
   AppendTo(&frames, WebSocketFrameHeader::kOpCodeContinuation, kFinal, "lo");
-  WriteFramesStub stub(OK);
+  predictor_->AddFramesToBeInput(frames);
+  WriteFramesStub stub(predictor_, OK);
   CompletionCallback callback;
 
   {
@@ -871,7 +991,8 @@ TEST_F(WebSocketDeflateStreamTest, WriteControlFrameBetweenDataFrames) {
 TEST_F(WebSocketDeflateStreamTest, WriteEmptyMessage) {
   ScopedVector<WebSocketFrame> frames;
   AppendTo(&frames, WebSocketFrameHeader::kOpCodeText, kFinal);
-  WriteFramesStub stub(OK);
+  predictor_->AddFramesToBeInput(frames);
+  WriteFramesStub stub(predictor_, OK);
   CompletionCallback callback;
 
   {
@@ -888,10 +1009,39 @@ TEST_F(WebSocketDeflateStreamTest, WriteEmptyMessage) {
   EXPECT_EQ(std::string("\x02\x00", 2), ToString(frames_passed[0]));
 }
 
+TEST_F(WebSocketDeflateStreamTest, WriteUncompressedMessage) {
+  ScopedVector<WebSocketFrame> frames;
+  AppendTo(&frames, WebSocketFrameHeader::kOpCodeText, kNoFlag, "AAAA");
+  AppendTo(&frames, WebSocketFrameHeader::kOpCodeContinuation, kFinal, "AAA");
+  predictor_->AddFramesToBeInput(frames);
+  WriteFramesStub stub(predictor_, OK);
+  CompletionCallback callback;
+
+  predictor_->set_result(WebSocketDeflatePredictor::DO_NOT_DEFLATE);
+
+  {
+    InSequence s;
+    EXPECT_CALL(*mock_stream_, WriteFrames(&frames, _))
+        .WillOnce(Invoke(&stub, &WriteFramesStub::Call));
+  }
+  ASSERT_EQ(OK, deflate_stream_->WriteFrames(&frames, callback));
+  const ScopedVector<WebSocketFrame>& frames_passed = *stub.frames();
+  ASSERT_EQ(2u, frames_passed.size());
+  EXPECT_EQ(WebSocketFrameHeader::kOpCodeText, frames_passed[0]->header.opcode);
+  EXPECT_FALSE(frames_passed[0]->header.final);
+  EXPECT_FALSE(frames_passed[0]->header.reserved1);
+  EXPECT_EQ("AAAA", ToString(frames_passed[0]));
+  EXPECT_EQ(WebSocketFrameHeader::kOpCodeContinuation,
+            frames_passed[1]->header.opcode);
+  EXPECT_TRUE(frames_passed[1]->header.final);
+  EXPECT_FALSE(frames_passed[1]->header.reserved1);
+  EXPECT_EQ("AAA", ToString(frames_passed[1]));
+}
+
 TEST_F(WebSocketDeflateStreamTest, LargeDeflatedFramesShouldBeSplit) {
   WebSocketDeflater deflater(WebSocketDeflater::TAKE_OVER_CONTEXT);
   LinearCongruentialGenerator lcg(133);
-  WriteFramesStub stub(OK);
+  WriteFramesStub stub(predictor_, OK);
   CompletionCallback callback;
   const size_t size = 1024;
 
@@ -912,9 +1062,11 @@ TEST_F(WebSocketDeflateStreamTest, LargeDeflatedFramesShouldBeSplit) {
     deflater.AddBytes(data.data(), data.size());
     FrameFlag flag = is_final ? kFinal : kNoFlag;
     AppendTo(&frames, WebSocketFrameHeader::kOpCodeBinary, flag, data);
+    predictor_->AddFramesToBeInput(frames);
     ASSERT_EQ(OK, deflate_stream_->WriteFrames(&frames, callback));
-    for (size_t i = 0; i < stub.frames()->size(); ++i)
-      total_compressed_frames.push_back((*stub.frames())[i]);
+    total_compressed_frames.insert(total_compressed_frames.end(),
+                                   stub.frames()->begin(),
+                                   stub.frames()->end());
     stub.frames()->weak_clear();
     if (is_final)
       break;
@@ -945,7 +1097,8 @@ TEST_F(WebSocketDeflateStreamTest, WriteMultipleMessages) {
   ScopedVector<WebSocketFrame> frames;
   AppendTo(&frames, WebSocketFrameHeader::kOpCodeText, kFinal, "Hello");
   AppendTo(&frames, WebSocketFrameHeader::kOpCodeText, kFinal, "Hello");
-  WriteFramesStub stub(OK);
+  predictor_->AddFramesToBeInput(frames);
+  WriteFramesStub stub(predictor_, OK);
   CompletionCallback callback;
 
   {
@@ -972,7 +1125,8 @@ TEST_F(WebSocketDeflateStreamWithDoNotTakeOverContextTest,
   ScopedVector<WebSocketFrame> frames;
   AppendTo(&frames, WebSocketFrameHeader::kOpCodeText, kFinal, "Hello");
   AppendTo(&frames, WebSocketFrameHeader::kOpCodeText, kFinal, "Hello");
-  WriteFramesStub stub(OK);
+  predictor_->AddFramesToBeInput(frames);
+  WriteFramesStub stub(predictor_, OK);
   CompletionCallback callback;
 
   {
@@ -995,6 +1149,58 @@ TEST_F(WebSocketDeflateStreamWithDoNotTakeOverContextTest,
             ToString(frames_passed[1]));
 }
 
+// In order to check the stream works correctly for multiple
+// "PossiblyCompressedMessage"s, we test various messages at one test case.
+TEST_F(WebSocketDeflateStreamWithDoNotTakeOverContextTest,
+       WritePossiblyCompressMessages) {
+  ScopedVector<WebSocketFrame> frames;
+  AppendTo(&frames, WebSocketFrameHeader::kOpCodeText, kNoFlag, "He");
+  AppendTo(&frames, WebSocketFrameHeader::kOpCodeContinuation, kFinal, "llo");
+  AppendTo(&frames, WebSocketFrameHeader::kOpCodeText, kNoFlag, "AAAAAAAAAA");
+  AppendTo(&frames, WebSocketFrameHeader::kOpCodeContinuation, kFinal, "AA");
+  AppendTo(&frames, WebSocketFrameHeader::kOpCodeText, kNoFlag, "XX");
+  AppendTo(&frames, WebSocketFrameHeader::kOpCodeContinuation, kFinal, "YY");
+  predictor_->AddFramesToBeInput(frames);
+  WriteFramesStub stub(predictor_, OK);
+  CompletionCallback callback;
+  predictor_->set_result(WebSocketDeflatePredictor::TRY_DEFLATE);
+
+  {
+    InSequence s;
+    EXPECT_CALL(*mock_stream_, WriteFrames(&frames, _))
+        .WillOnce(Invoke(&stub, &WriteFramesStub::Call));
+  }
+  ASSERT_EQ(OK, deflate_stream_->WriteFrames(&frames, callback));
+  const ScopedVector<WebSocketFrame>& frames_passed = *stub.frames();
+  ASSERT_EQ(5u, frames_passed.size());
+
+  EXPECT_EQ(WebSocketFrameHeader::kOpCodeText, frames_passed[0]->header.opcode);
+  EXPECT_FALSE(frames_passed[0]->header.final);
+  EXPECT_FALSE(frames_passed[0]->header.reserved1);
+  EXPECT_EQ("He", ToString(frames_passed[0]));
+  EXPECT_EQ(WebSocketFrameHeader::kOpCodeContinuation,
+            frames_passed[1]->header.opcode);
+  EXPECT_TRUE(frames_passed[1]->header.final);
+  EXPECT_FALSE(frames_passed[1]->header.reserved1);
+  EXPECT_EQ("llo", ToString(frames_passed[1]));
+
+  EXPECT_EQ(WebSocketFrameHeader::kOpCodeText, frames_passed[2]->header.opcode);
+  EXPECT_TRUE(frames_passed[2]->header.final);
+  EXPECT_TRUE(frames_passed[2]->header.reserved1);
+  EXPECT_EQ(std::string("\x72\x74\x44\x00\x00\x00", 6),
+            ToString(frames_passed[2]));
+
+  EXPECT_EQ(WebSocketFrameHeader::kOpCodeText, frames_passed[3]->header.opcode);
+  EXPECT_FALSE(frames_passed[3]->header.final);
+  EXPECT_FALSE(frames_passed[3]->header.reserved1);
+  EXPECT_EQ("XX", ToString(frames_passed[3]));
+  EXPECT_EQ(WebSocketFrameHeader::kOpCodeContinuation,
+            frames_passed[4]->header.opcode);
+  EXPECT_TRUE(frames_passed[4]->header.final);
+  EXPECT_FALSE(frames_passed[4]->header.reserved1);
+  EXPECT_EQ("YY", ToString(frames_passed[4]));
+}
+
 }  // namespace
 
 }  // namespace net