Limit the total memory usage for Stream instances

content/browser/byte_stream.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 "content/browser/byte_stream.h"
 
 #include <deque>
 #include <set>
 #include <utility>
 
@@ skipping 43 matching lines @@
   void SetPeer(ByteStreamReaderImpl* peer,
                scoped_refptr<base::SequencedTaskRunner> peer_task_runner,
                scoped_refptr<LifetimeFlag> peer_lifetime_flag);
 
   // Overridden from ByteStreamWriter.
   virtual bool Write(scoped_refptr<net::IOBuffer> buffer,
                      size_t byte_count) OVERRIDE;
   virtual void Flush() OVERRIDE;
   virtual void Close(int status) OVERRIDE;
   virtual void RegisterCallback(const base::Closure& source_callback) OVERRIDE;
+  virtual size_t TotalBufferedBytes() const OVERRIDE;
 
   // PostTask target from |ByteStreamReaderImpl::MaybeUpdateInput|.
   static void UpdateWindow(scoped_refptr<LifetimeFlag> lifetime_flag,
                            ByteStreamWriterImpl* target,
                            size_t bytes_consumed);
 
  private:
   // Called from UpdateWindow when object existence has been validated.
   void UpdateWindowInternal(size_t bytes_consumed);
 
@@ skipping 112 matching lines @@
       my_task_runner_(task_runner),
       my_lifetime_flag_(lifetime_flag),
       input_contents_size_(0),
       output_size_used_(0),
       peer_(NULL) {
   DCHECK(my_lifetime_flag_.get());
   my_lifetime_flag_->is_alive = true;
 }
 
 ByteStreamWriterImpl::~ByteStreamWriterImpl() {
+  DCHECK(my_task_runner_->RunsTasksOnCurrentThread());
   my_lifetime_flag_->is_alive = false;
 }
 
 void ByteStreamWriterImpl::SetPeer(
     ByteStreamReaderImpl* peer,
     scoped_refptr<base::SequencedTaskRunner> peer_task_runner,
     scoped_refptr<LifetimeFlag> peer_lifetime_flag) {
   peer_ = peer;
   peer_task_runner_ = peer_task_runner;
   peer_lifetime_flag_ = peer_lifetime_flag;
 }
 
 bool ByteStreamWriterImpl::Write(
     scoped_refptr<net::IOBuffer> buffer, size_t byte_count) {
   DCHECK(my_task_runner_->RunsTasksOnCurrentThread());
 
   input_contents_.push_back(std::make_pair(buffer, byte_count));
   input_contents_size_ += byte_count;
 
   // Arbitrarily, we buffer to a third of the total size before sending.
   if (input_contents_size_ > total_buffer_size_ / kFractionBufferBeforeSending)
     PostToPeer(false, 0);
 
-  return (input_contents_size_ + output_size_used_ <= total_buffer_size_);
+  return (TotalBufferedBytes() <= total_buffer_size_);
 }
 
 void ByteStreamWriterImpl::Flush() {
   DCHECK(my_task_runner_->RunsTasksOnCurrentThread());
   if (input_contents_size_ > 0)
     PostToPeer(false, 0);
 }
 
 void ByteStreamWriterImpl::Close(int status) {
   DCHECK(my_task_runner_->RunsTasksOnCurrentThread());
   PostToPeer(true, status);
 }
 
 void ByteStreamWriterImpl::RegisterCallback(
     const base::Closure& source_callback) {
   DCHECK(my_task_runner_->RunsTasksOnCurrentThread());
   space_available_callback_ = source_callback;
 }
 
+size_t ByteStreamWriterImpl::TotalBufferedBytes() const {
+  DCHECK(my_task_runner_->RunsTasksOnCurrentThread());
+  return input_contents_size_ + output_size_used_;

[kinuko, 2013/08/14 09:32:23]

Could this overflow?

[tyoshino (SeeGerritForStatus), 2013/08/15 04:15:33]

Added overflow check in Write()

+}
+
 // static
 void ByteStreamWriterImpl::UpdateWindow(
     scoped_refptr<LifetimeFlag> lifetime_flag, ByteStreamWriterImpl* target,
     size_t bytes_consumed) {
   // If the target object isn't alive anymore, we do nothing.
   if (!lifetime_flag->is_alive) return;
 
   target->UpdateWindowInternal(bytes_consumed);
 }
 
@@ skipping 46 matching lines @@
       my_lifetime_flag_(lifetime_flag),
       received_status_(false),
       status_(0),
       unreported_consumed_bytes_(0),
       peer_(NULL) {
   DCHECK(my_lifetime_flag_.get());
   my_lifetime_flag_->is_alive = true;
 }
 
 ByteStreamReaderImpl::~ByteStreamReaderImpl() {
+  DCHECK(my_task_runner_->RunsTasksOnCurrentThread());
   my_lifetime_flag_->is_alive = false;
 }
 
 void ByteStreamReaderImpl::SetPeer(
     ByteStreamWriterImpl* peer,
     scoped_refptr<base::SequencedTaskRunner> peer_task_runner,
     scoped_refptr<LifetimeFlag> peer_lifetime_flag) {
   peer_ = peer;
   peer_task_runner_ = peer_task_runner;
   peer_lifetime_flag_ = peer_lifetime_flag;
@@ skipping 117 matching lines @@
   ByteStreamReaderImpl* out = new ByteStreamReaderImpl(
       output_task_runner, output_flag, buffer_size);
 
   in->SetPeer(out, output_task_runner, output_flag);
   out->SetPeer(in, input_task_runner, input_flag);
   input->reset(in);
   output->reset(out);
 }
 
 }  // namespace content