Fix MessageReader not to invoke message callback from done callback.

remoting/protocol/message_reader.cc

 // Copyright (c) 2012 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 "remoting/protocol/message_reader.h"
 
 #include "base/bind.h"
 #include "base/callback.h"
 #include "base/compiler_specific.h"
 #include "base/location.h"
 #include "base/thread_task_runner_handle.h"
+#include "base/single_thread_task_runner.h"
 #include "net/base/io_buffer.h"
 #include "net/base/net_errors.h"
 #include "net/socket/socket.h"
 #include "remoting/base/compound_buffer.h"
 #include "remoting/proto/internal.pb.h"
 
 namespace remoting {
 namespace protocol {
 
 static const int kReadBufferSize = 4096;
 
 MessageReader::MessageReader()
     : socket_(NULL),
       read_pending_(false),
       pending_messages_(0),
       closed_(false),
       ALLOW_THIS_IN_INITIALIZER_LIST(weak_factory_(this)) {
 }
 
 void MessageReader::Init(net::Socket* socket,
                          const MessageReceivedCallback& callback) {
   DCHECK(CalledOnValidThread());
   message_received_callback_ = callback;
   DCHECK(socket);
   socket_ = socket;
-  DoRead();
+  DoRead(false);
 }
 
 MessageReader::~MessageReader() {
 }
 
-void MessageReader::DoRead() {
+void MessageReader::DoRead(bool in_callback) {
   DCHECK(CalledOnValidThread());
   // Don't try to read again if there is another read pending or we
   // have messages that we haven't finished processing yet.
   while (!closed_ && !read_pending_ && pending_messages_ == 0) {
     read_buffer_ = new net::IOBuffer(kReadBufferSize);
     int result = socket_->Read(
         read_buffer_, kReadBufferSize,
         base::Bind(&MessageReader::OnRead, weak_factory_.GetWeakPtr()));
-    HandleReadResult(result);
+    HandleReadResult(result, in_callback);
   }
 }
 
 void MessageReader::OnRead(int result) {
   DCHECK(CalledOnValidThread());
   DCHECK(read_pending_);
   read_pending_ = false;
 
   if (!closed_) {
-    HandleReadResult(result);
-    DoRead();
+    HandleReadResult(result, false);
+    DoRead(false);
   }
 }
 
-void MessageReader::HandleReadResult(int result) {
+void MessageReader::HandleReadResult(int result, bool in_callback) {
   DCHECK(CalledOnValidThread());
   if (closed_)
     return;
 
   if (result > 0) {
-    OnDataReceived(read_buffer_, result);
+    OnDataReceived(read_buffer_, result, in_callback);
   } else if (result == net::ERR_IO_PENDING) {
     read_pending_ = true;
   } else {
     if (result != net::ERR_CONNECTION_CLOSED) {
       LOG(ERROR) << "Read() returned error " << result;
     }
     // Stop reading after any error.
     closed_ = true;
   }
 }
 
-void MessageReader::OnDataReceived(net::IOBuffer* data, int data_size) {
+void MessageReader::OnDataReceived(net::IOBuffer* data, int data_size,
+                                   bool in_callback) {
   DCHECK(CalledOnValidThread());
   message_decoder_.AddData(data, data_size);
 
   // Get list of all new messages first, and then call the callback
   // for all of them.
   std::vector<CompoundBuffer*> new_messages;
   while (true) {
     CompoundBuffer* buffer = message_decoder_.GetNextMessage();
     if (!buffer)
       break;
     new_messages.push_back(buffer);
   }
 
   pending_messages_ += new_messages.size();
 
   for (std::vector<CompoundBuffer*>::iterator it = new_messages.begin();
        it != new_messages.end(); ++it) {
-    message_received_callback_.Run(
-        scoped_ptr<CompoundBuffer>(*it),
-        base::Bind(&MessageReader::OnMessageDone,
-                   weak_factory_.GetWeakPtr()));
+    scoped_ptr<CompoundBuffer> message(*it);
+    if (in_callback) {
+      base::ThreadTaskRunnerHandle::Get()->PostTask(
+          FROM_HERE, base::Bind(&MessageReader::RunCallback,
+                                weak_factory_.GetWeakPtr(),
+                                base::Passed(&message)));
+    } else {
+      RunCallback(message.Pass());
+    }
   }
 }
 
+void MessageReader::RunCallback(scoped_ptr<CompoundBuffer> message) {
+  message_received_callback_.Run(
+      message.Pass(), base::Bind(&MessageReader::OnMessageDone,
+                                 weak_factory_.GetWeakPtr()));
+}
+
 void MessageReader::OnMessageDone() {
   DCHECK(CalledOnValidThread());
   pending_messages_--;
   DCHECK_GE(pending_messages_, 0);
 
+  // Start next read if necessary.
   if (!read_pending_)

[Wez, 2012/09/13 01:00:09]

nit: DoRead() already checks |read_pending_| so you don't need to check it here.

[Sergey Ulanov, 2012/09/13 01:40:21]

Done.

-    DoRead();  // Start next read if necessary.
+    DoRead(true);

[Wez, 2012/09/13 01:00:09]

Wouldn't it be simpler to have a member that you set in RunCallback around the
|message_received_callback_.Run()|, and which you clear afterwards, that
prevents the DoRead() from within OnMessageDone()?

[Sergey Ulanov, 2012/09/13 01:40:21]

I actually noticed another bug in this code - it didn't handle properly the case
when MessageReader is deleted when running OnMessageCallback. So I changed it to
always post a task for new messages.
Simplified the code significantly.

 }
 
 }  // namespace protocol
 }  // namespace remoting