Implment BytesConsumer::tee

third_party/WebKit/Source/modules/fetch/BytesConsumerForDataConsumerHandle.cpp

Index: third_party/WebKit/Source/modules/fetch/BytesConsumerForDataConsumerHandle.cpp
diff --git a/third_party/WebKit/Source/modules/fetch/BytesConsumerForDataConsumerHandle.cpp b/third_party/WebKit/Source/modules/fetch/BytesConsumerForDataConsumerHandle.cpp
index f77611c5b857f3d687c58aac029616abc321e42f..e4a61abfc1dd129f84f3eca9c18c0f81471d738f 100644
--- a/third_party/WebKit/Source/modules/fetch/BytesConsumerForDataConsumerHandle.cpp
+++ b/third_party/WebKit/Source/modules/fetch/BytesConsumerForDataConsumerHandle.cpp
@@ -4,20 +4,30 @@
 
 #include "modules/fetch/BytesConsumerForDataConsumerHandle.h"
 
+#include "core/dom/ExecutionContext.h"
+#include "core/dom/TaskRunnerHelper.h"
+#include "public/platform/WebTaskRunner.h"
+#include "public/platform/WebTraceLocation.h"
+#include "wtf/Functional.h"
+
 #include <algorithm>
 #include <string.h>
 
 namespace blink {
 
-BytesConsumerForDataConsumerHandle::BytesConsumerForDataConsumerHandle(std::unique_ptr<FetchDataConsumerHandle> handle)
-    : m_reader(handle->obtainFetchDataReader(this))
+BytesConsumerForDataConsumerHandle::BytesConsumerForDataConsumerHandle(ExecutionContext* executionContext, std::unique_ptr<FetchDataConsumerHandle> handle)
+    : m_executionContext(executionContext)
+    , m_reader(handle->obtainFetchDataReader(this))
 {
 }
 
-BytesConsumerForDataConsumerHandle::~BytesConsumerForDataConsumerHandle() {}
+BytesConsumerForDataConsumerHandle::~BytesConsumerForDataConsumerHandle()
+{
+}
 
 BytesConsumer::Result BytesConsumerForDataConsumerHandle::read(char* buffer, size_t size, size_t* readSize)
 {
+    DCHECK(!m_isInTwoPhaseRead);
     *readSize = 0;
     if (m_state == InternalState::Closed)
         return Result::Done;
@@ -45,6 +55,7 @@ BytesConsumer::Result BytesConsumerForDataConsumerHandle::read(char* buffer, siz
 
 BytesConsumer::Result BytesConsumerForDataConsumerHandle::beginRead(const char** buffer, size_t* available)
 {
+    DCHECK(!m_isInTwoPhaseRead);
     *buffer = nullptr;
     *available = 0;
     if (m_state == InternalState::Closed)
@@ -55,6 +66,7 @@ BytesConsumer::Result BytesConsumerForDataConsumerHandle::beginRead(const char**
     WebDataConsumerHandle::Result r = m_reader->beginRead(reinterpret_cast<const void**>(buffer), WebDataConsumerHandle::FlagNone, available);
     switch (r) {
     case WebDataConsumerHandle::Ok:
+        m_isInTwoPhaseRead = true;
         return Result::Ok;
     case WebDataConsumerHandle::ShouldWait:
         return Result::ShouldWait;
@@ -73,12 +85,20 @@ BytesConsumer::Result BytesConsumerForDataConsumerHandle::beginRead(const char**
 
 BytesConsumer::Result BytesConsumerForDataConsumerHandle::endRead(size_t read)
 {
+    DCHECK(m_isInTwoPhaseRead);
+    m_isInTwoPhaseRead = false;
     DCHECK(m_state == InternalState::Readable || m_state == InternalState::Waiting);
     WebDataConsumerHandle::Result r = m_reader->endRead(read);
-    if (r == WebDataConsumerHandle::Ok)
-        return Result::Ok;
-    error();
-    return Result::Error;
+    if (r != WebDataConsumerHandle::Ok) {
+        m_hasPendingNotification = false;
+        error();
+        return Result::Error;
+    }
+    if (m_hasPendingNotification) {
+        m_hasPendingNotification = false;
+        TaskRunnerHelper::get(TaskType::Networking, m_executionContext)->postTask(BLINK_FROM_HERE, WTF::bind(&BytesConsumerForDataConsumerHandle::notify, wrapPersistent(this)));
+    }
+    return Result::Ok;
 }
 
 PassRefPtr<BlobDataHandle> BytesConsumerForDataConsumerHandle::drainAsBlobDataHandle(BlobSizePolicy policy)
@@ -113,17 +133,18 @@ void BytesConsumerForDataConsumerHandle::setClient(BytesConsumer::Client* client
 {
     DCHECK(!m_client);
     DCHECK(client);
-    m_client = client;
+    if (m_state == InternalState::Readable || m_state == InternalState::Waiting)
+        m_client = client;
 }
 
 void BytesConsumerForDataConsumerHandle::clearClient()
 {
-    DCHECK(m_client);
     m_client = nullptr;
 }
 
 void BytesConsumerForDataConsumerHandle::cancel()
 {
+    DCHECK(!m_isInTwoPhaseRead);
     if (m_state == InternalState::Readable || m_state == InternalState::Waiting) {
         // We don't want the client to be notified in this case.
         BytesConsumer::Client* client = m_client;
@@ -141,26 +162,31 @@ BytesConsumer::PublicState BytesConsumerForDataConsumerHandle::getPublicState()
 void BytesConsumerForDataConsumerHandle::didGetReadable()
 {
     DCHECK(m_state == InternalState::Readable || m_state == InternalState::Waiting);
+    if (m_isInTwoPhaseRead) {
+        m_hasPendingNotification = true;
+        return;
+    }
     // Perform zero-length read to call check handle's status.
     size_t readSize;
     WebDataConsumerHandle::Result result = m_reader->read(nullptr, 0, WebDataConsumerHandle::FlagNone, &readSize);
+    BytesConsumer::Client* client = m_client;
     switch (result) {
     case WebDataConsumerHandle::Ok:
     case WebDataConsumerHandle::ShouldWait:
-        if (m_client)
-            m_client->onStateChange();
+        if (client)
+            client->onStateChange();
         return;
     case WebDataConsumerHandle::Done:
         close();
-        if (m_client)
-            m_client->onStateChange();
+        if (client)
+            client->onStateChange();
         return;
     case WebDataConsumerHandle::Busy:
     case WebDataConsumerHandle::ResourceExhausted:
     case WebDataConsumerHandle::UnexpectedError:
         error();
-        if (m_client)
-            m_client->onStateChange();
+        if (client)
+            client->onStateChange();
         return;
     }
     return;
@@ -168,27 +194,39 @@ void BytesConsumerForDataConsumerHandle::didGetReadable()
 
 DEFINE_TRACE(BytesConsumerForDataConsumerHandle)
 {
+    visitor->trace(m_executionContext);
     visitor->trace(m_client);
     BytesConsumer::trace(visitor);
 }
 
 void BytesConsumerForDataConsumerHandle::close()
 {
+    DCHECK(!m_isInTwoPhaseRead);
     if (m_state == InternalState::Closed)
         return;
     DCHECK(m_state == InternalState::Readable || m_state == InternalState::Waiting);
     m_state = InternalState::Closed;
     m_reader = nullptr;
+    clearClient();
 }
 
 void BytesConsumerForDataConsumerHandle::error()
 {
+    DCHECK(!m_isInTwoPhaseRead);
     if (m_state == InternalState::Errored)
         return;
     DCHECK(m_state == InternalState::Readable || m_state == InternalState::Waiting);
     m_state = InternalState::Errored;
     m_reader = nullptr;
     m_error = Error("error");
+    clearClient();
+}
+
+void BytesConsumerForDataConsumerHandle::notify()
+{
+    if (m_state == InternalState::Closed || m_state == InternalState::Errored)
+        return;
+    didGetReadable();
 }
 
 } // namespace blink