// Copyright 2014 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/quic/quic_flow_controller.h"

#include <memory>

#include "base/format_macros.h"
#include "base/strings/stringprintf.h"
#include "net/quic/quic_flags.h"
#include "net/quic/quic_utils.h"
#include "net/quic/test_tools/quic_connection_peer.h"
#include "net/quic/test_tools/quic_flow_controller_peer.h"
#include "net/quic/test_tools/quic_sent_packet_manager_peer.h"
#include "net/quic/test_tools/quic_test_utils.h"
#include "net/test/gtest_util.h"
#include "testing/gmock/include/gmock/gmock.h"

using testing::_;

namespace net {
namespace test {

// Receive window auto-tuning uses RTT in its logic.
const int64_t kRtt = 100;

class QuicFlowControllerTest : public ::testing::Test {
 public:
  QuicFlowControllerTest()
      : stream_id_(1234),
        send_window_(kInitialSessionFlowControlWindowForTest),
        receive_window_(kInitialSessionFlowControlWindowForTest),
        connection_(&helper_, &alarm_factory_, Perspective::IS_CLIENT) {}

  void Initialize() {
    flow_controller_.reset(
        new QuicFlowController(&connection_, stream_id_, Perspective::IS_CLIENT,
                               send_window_, receive_window_, false));
  }

 protected:
  QuicStreamId stream_id_;
  QuicByteCount send_window_;
  QuicByteCount receive_window_;
  std::unique_ptr<QuicFlowController> flow_controller_;
  MockQuicConnectionHelper helper_;
  MockAlarmFactory alarm_factory_;
  MockQuicConnection connection_;
};

TEST_F(QuicFlowControllerTest, SendingBytes) {
  Initialize();

  EXPECT_FALSE(flow_controller_->IsBlocked());
  EXPECT_FALSE(flow_controller_->FlowControlViolation());
  EXPECT_EQ(send_window_, flow_controller_->SendWindowSize());

  // Send some bytes, but not enough to block.
  flow_controller_->AddBytesSent(send_window_ / 2);
  EXPECT_FALSE(flow_controller_->IsBlocked());
  EXPECT_EQ(send_window_ / 2, flow_controller_->SendWindowSize());

  // Send enough bytes to block.
  flow_controller_->AddBytesSent(send_window_ / 2);
  EXPECT_TRUE(flow_controller_->IsBlocked());
  EXPECT_EQ(0u, flow_controller_->SendWindowSize());

  // BLOCKED frame should get sent.
  EXPECT_CALL(connection_, SendBlocked(stream_id_)).Times(1);
  flow_controller_->MaybeSendBlocked();

  // Update the send window, and verify this has unblocked.
  EXPECT_TRUE(flow_controller_->UpdateSendWindowOffset(2 * send_window_));
  EXPECT_FALSE(flow_controller_->IsBlocked());
  EXPECT_EQ(send_window_, flow_controller_->SendWindowSize());

  // Updating with a smaller offset doesn't change anything.
  EXPECT_FALSE(flow_controller_->UpdateSendWindowOffset(send_window_ / 10));
  EXPECT_EQ(send_window_, flow_controller_->SendWindowSize());

  // Try to send more bytes, violating flow control.
  EXPECT_CALL(connection_,
              CloseConnection(QUIC_FLOW_CONTROL_SENT_TOO_MUCH_DATA, _, _));
  EXPECT_DFATAL(flow_controller_->AddBytesSent(send_window_ * 10),
                base::StringPrintf("Trying to send an extra %" PRIu64 " bytes",
                                   send_window_ * 10));
  EXPECT_TRUE(flow_controller_->IsBlocked());
  EXPECT_EQ(0u, flow_controller_->SendWindowSize());
}

TEST_F(QuicFlowControllerTest, ReceivingBytes) {
  Initialize();

  EXPECT_FALSE(flow_controller_->IsBlocked());
  EXPECT_FALSE(flow_controller_->FlowControlViolation());
  EXPECT_EQ(kInitialSessionFlowControlWindowForTest,
            QuicFlowControllerPeer::ReceiveWindowSize(flow_controller_.get()));

  // Receive some bytes, updating highest received offset, but not enough to
  // fill flow control receive window.
  EXPECT_TRUE(
      flow_controller_->UpdateHighestReceivedOffset(1 + receive_window_ / 2));
  EXPECT_FALSE(flow_controller_->FlowControlViolation());
  EXPECT_EQ((receive_window_ / 2) - 1,
            QuicFlowControllerPeer::ReceiveWindowSize(flow_controller_.get()));

  // Consume enough bytes to send a WINDOW_UPDATE frame.
  EXPECT_CALL(connection_, SendWindowUpdate(stream_id_, ::testing::_)).Times(1);

  flow_controller_->AddBytesConsumed(1 + receive_window_ / 2);

  // Result is that once again we have a fully open receive window.
  EXPECT_FALSE(flow_controller_->FlowControlViolation());
  EXPECT_EQ(kInitialSessionFlowControlWindowForTest,
            QuicFlowControllerPeer::ReceiveWindowSize(flow_controller_.get()));
}

TEST_F(QuicFlowControllerTest, OnlySendBlockedFrameOncePerOffset) {
  Initialize();

  // Test that we don't send duplicate BLOCKED frames. We should only send one
  // BLOCKED frame at a given send window offset.
  EXPECT_FALSE(flow_controller_->IsBlocked());
  EXPECT_FALSE(flow_controller_->FlowControlViolation());
  EXPECT_EQ(send_window_, flow_controller_->SendWindowSize());

  // Send enough bytes to block.
  flow_controller_->AddBytesSent(send_window_);
  EXPECT_TRUE(flow_controller_->IsBlocked());
  EXPECT_EQ(0u, flow_controller_->SendWindowSize());

  // Expect that 2 BLOCKED frames should get sent in total.
  EXPECT_CALL(connection_, SendBlocked(stream_id_)).Times(2);

  // BLOCKED frame should get sent.
  flow_controller_->MaybeSendBlocked();

  // BLOCKED frame should not get sent again until our send offset changes.
  flow_controller_->MaybeSendBlocked();
  flow_controller_->MaybeSendBlocked();
  flow_controller_->MaybeSendBlocked();
  flow_controller_->MaybeSendBlocked();
  flow_controller_->MaybeSendBlocked();

  // Update the send window, then send enough bytes to block again.
  EXPECT_TRUE(flow_controller_->UpdateSendWindowOffset(2 * send_window_));
  EXPECT_FALSE(flow_controller_->IsBlocked());
  EXPECT_EQ(send_window_, flow_controller_->SendWindowSize());
  flow_controller_->AddBytesSent(send_window_);
  EXPECT_TRUE(flow_controller_->IsBlocked());
  EXPECT_EQ(0u, flow_controller_->SendWindowSize());

  // BLOCKED frame should get sent as send offset has changed.
  flow_controller_->MaybeSendBlocked();
}

TEST_F(QuicFlowControllerTest, ReceivingBytesFastIncreasesFlowWindow) {
  FLAGS_quic_auto_tune_receive_window = true;
  // This test will generate two WINDOW_UPDATE frames.
  EXPECT_CALL(connection_, SendWindowUpdate(stream_id_, ::testing::_)).Times(2);

  Initialize();
  flow_controller_->set_auto_tune_receive_window(true);

  // Make sure clock is inititialized.
  connection_.AdvanceTime(QuicTime::Delta::FromMilliseconds(1));

  QuicSentPacketManagerInterface* manager =
      QuicConnectionPeer::GetSentPacketManager(&connection_);

  RttStats* rtt_stats = const_cast<RttStats*>(manager->GetRttStats());
  rtt_stats->UpdateRtt(QuicTime::Delta::FromMilliseconds(kRtt),
                       QuicTime::Delta::Zero(), QuicTime::Zero());

  EXPECT_FALSE(flow_controller_->IsBlocked());
  EXPECT_FALSE(flow_controller_->FlowControlViolation());
  EXPECT_EQ(kInitialSessionFlowControlWindowForTest,
            QuicFlowControllerPeer::ReceiveWindowSize(flow_controller_.get()));

  QuicByteCount threshold =
      QuicFlowControllerPeer::WindowUpdateThreshold(flow_controller_.get());

  QuicStreamOffset receive_offset = threshold + 1;
  // Receive some bytes, updating highest received offset, but not enough to
  // fill flow control receive window.
  EXPECT_TRUE(flow_controller_->UpdateHighestReceivedOffset(receive_offset));
  EXPECT_FALSE(flow_controller_->FlowControlViolation());
  EXPECT_EQ(kInitialSessionFlowControlWindowForTest - receive_offset,
            QuicFlowControllerPeer::ReceiveWindowSize(flow_controller_.get()));

  // Consume enough bytes to send a WINDOW_UPDATE frame.
  flow_controller_->AddBytesConsumed(threshold + 1);
  // Result is that once again we have a fully open receive window.
  EXPECT_FALSE(flow_controller_->FlowControlViolation());
  EXPECT_EQ(kInitialSessionFlowControlWindowForTest,
            QuicFlowControllerPeer::ReceiveWindowSize(flow_controller_.get()));

  // Move time forward, but by less than two RTTs.  Then receive and consume
  // some more, forcing a second WINDOW_UPDATE with an increased max window
  // size.
  connection_.AdvanceTime(QuicTime::Delta::FromMilliseconds(2 * kRtt - 1));
  receive_offset += threshold + 1;
  EXPECT_TRUE(flow_controller_->UpdateHighestReceivedOffset(receive_offset));
  flow_controller_->AddBytesConsumed(threshold + 1);
  EXPECT_FALSE(flow_controller_->FlowControlViolation());
  QuicByteCount new_threshold =
      QuicFlowControllerPeer::WindowUpdateThreshold(flow_controller_.get());
  EXPECT_GT(new_threshold, threshold);
}

TEST_F(QuicFlowControllerTest, ReceivingBytesFastStatusQuo) {
  FLAGS_quic_auto_tune_receive_window = false;
  // This test will generate two WINDOW_UPDATE frames.
  EXPECT_CALL(connection_, SendWindowUpdate(stream_id_, ::testing::_)).Times(2);

  Initialize();
  flow_controller_->set_auto_tune_receive_window(true);

  // Make sure clock is inititialized.
  connection_.AdvanceTime(QuicTime::Delta::FromMilliseconds(1));

  QuicSentPacketManagerInterface* manager =
      QuicConnectionPeer::GetSentPacketManager(&connection_);

  RttStats* rtt_stats = const_cast<RttStats*>(manager->GetRttStats());
  rtt_stats->UpdateRtt(QuicTime::Delta::FromMilliseconds(kRtt),
                       QuicTime::Delta::Zero(), QuicTime::Zero());

  EXPECT_FALSE(flow_controller_->IsBlocked());
  EXPECT_FALSE(flow_controller_->FlowControlViolation());
  EXPECT_EQ(kInitialSessionFlowControlWindowForTest,
            QuicFlowControllerPeer::ReceiveWindowSize(flow_controller_.get()));

  QuicByteCount threshold =
      QuicFlowControllerPeer::WindowUpdateThreshold(flow_controller_.get());

  QuicStreamOffset receive_offset = threshold + 1;
  // Receive some bytes, updating highest received offset, but not enough to
  // fill flow control receive window.
  EXPECT_TRUE(flow_controller_->UpdateHighestReceivedOffset(receive_offset));
  EXPECT_FALSE(flow_controller_->FlowControlViolation());
  EXPECT_EQ(kInitialSessionFlowControlWindowForTest - receive_offset,
            QuicFlowControllerPeer::ReceiveWindowSize(flow_controller_.get()));

  // Consume enough bytes to send a WINDOW_UPDATE frame.
  flow_controller_->AddBytesConsumed(threshold + 1);
  // Result is that once again we have a fully open receive window.
  EXPECT_FALSE(flow_controller_->FlowControlViolation());
  EXPECT_EQ(kInitialSessionFlowControlWindowForTest,
            QuicFlowControllerPeer::ReceiveWindowSize(flow_controller_.get()));

  // Move time forward, but by less than two RTTs.  Then receive and consume
  // some more, forcing a second WINDOW_UPDATE with an increased max window
  // size.
  connection_.AdvanceTime(QuicTime::Delta::FromMilliseconds(2 * kRtt - 1));
  receive_offset += threshold + 1;
  EXPECT_TRUE(flow_controller_->UpdateHighestReceivedOffset(receive_offset));
  flow_controller_->AddBytesConsumed(threshold + 1);
  EXPECT_FALSE(flow_controller_->FlowControlViolation());
  QuicByteCount new_threshold =
      QuicFlowControllerPeer::WindowUpdateThreshold(flow_controller_.get());
  EXPECT_EQ(new_threshold, threshold);
}

TEST_F(QuicFlowControllerTest, ReceivingBytesNormalStableFlowWindow) {
  FLAGS_quic_auto_tune_receive_window = true;
  // This test will generate two WINDOW_UPDATE frames.
  EXPECT_CALL(connection_, SendWindowUpdate(stream_id_, ::testing::_)).Times(2);

  Initialize();
  flow_controller_->set_auto_tune_receive_window(true);

  // Make sure clock is inititialized.
  connection_.AdvanceTime(QuicTime::Delta::FromMilliseconds(1));

  QuicSentPacketManagerInterface* manager =
      QuicConnectionPeer::GetSentPacketManager(&connection_);
  RttStats* rtt_stats = const_cast<RttStats*>(manager->GetRttStats());
  rtt_stats->UpdateRtt(QuicTime::Delta::FromMilliseconds(kRtt),
                       QuicTime::Delta::Zero(), QuicTime::Zero());

  EXPECT_FALSE(flow_controller_->IsBlocked());
  EXPECT_FALSE(flow_controller_->FlowControlViolation());
  EXPECT_EQ(kInitialSessionFlowControlWindowForTest,
            QuicFlowControllerPeer::ReceiveWindowSize(flow_controller_.get()));

  QuicByteCount threshold =
      QuicFlowControllerPeer::WindowUpdateThreshold(flow_controller_.get());

  QuicStreamOffset receive_offset = threshold + 1;
  // Receive some bytes, updating highest received offset, but not enough to
  // fill flow control receive window.
  EXPECT_TRUE(flow_controller_->UpdateHighestReceivedOffset(receive_offset));
  EXPECT_FALSE(flow_controller_->FlowControlViolation());
  EXPECT_EQ(kInitialSessionFlowControlWindowForTest - receive_offset,
            QuicFlowControllerPeer::ReceiveWindowSize(flow_controller_.get()));

  flow_controller_->AddBytesConsumed(threshold + 1);

  // Result is that once again we have a fully open receive window.
  EXPECT_FALSE(flow_controller_->FlowControlViolation());
  EXPECT_EQ(kInitialSessionFlowControlWindowForTest,
            QuicFlowControllerPeer::ReceiveWindowSize(flow_controller_.get()));

  // Move time forward, but by more than two RTTs.  Then receive and consume
  // some more, forcing a second WINDOW_UPDATE with unchanged max window size.
  connection_.AdvanceTime(QuicTime::Delta::FromMilliseconds(2 * kRtt + 1));

  receive_offset += threshold + 1;
  EXPECT_TRUE(flow_controller_->UpdateHighestReceivedOffset(receive_offset));

  flow_controller_->AddBytesConsumed(threshold + 1);
  EXPECT_FALSE(flow_controller_->FlowControlViolation());

  QuicByteCount new_threshold =
      QuicFlowControllerPeer::WindowUpdateThreshold(flow_controller_.get());

  EXPECT_EQ(new_threshold, threshold);
}

TEST_F(QuicFlowControllerTest, ReceivingBytesNormalStatusQuo) {
  FLAGS_quic_auto_tune_receive_window = false;
  // This test will generate two WINDOW_UPDATE frames.
  EXPECT_CALL(connection_, SendWindowUpdate(stream_id_, ::testing::_)).Times(2);

  Initialize();
  flow_controller_->set_auto_tune_receive_window(true);

  // Make sure clock is inititialized.
  connection_.AdvanceTime(QuicTime::Delta::FromMilliseconds(1));

  QuicSentPacketManagerInterface* manager =
      QuicConnectionPeer::GetSentPacketManager(&connection_);
  RttStats* rtt_stats = const_cast<RttStats*>(manager->GetRttStats());
  rtt_stats->UpdateRtt(QuicTime::Delta::FromMilliseconds(kRtt),
                       QuicTime::Delta::Zero(), QuicTime::Zero());

  EXPECT_FALSE(flow_controller_->IsBlocked());
  EXPECT_FALSE(flow_controller_->FlowControlViolation());
  EXPECT_EQ(kInitialSessionFlowControlWindowForTest,
            QuicFlowControllerPeer::ReceiveWindowSize(flow_controller_.get()));

  QuicByteCount threshold =
      QuicFlowControllerPeer::WindowUpdateThreshold(flow_controller_.get());

  QuicStreamOffset receive_offset = threshold + 1;
  // Receive some bytes, updating highest received offset, but not enough to
  // fill flow control receive window.
  EXPECT_TRUE(flow_controller_->UpdateHighestReceivedOffset(receive_offset));
  EXPECT_FALSE(flow_controller_->FlowControlViolation());
  EXPECT_EQ(kInitialSessionFlowControlWindowForTest - receive_offset,
            QuicFlowControllerPeer::ReceiveWindowSize(flow_controller_.get()));

  flow_controller_->AddBytesConsumed(threshold + 1);

  // Result is that once again we have a fully open receive window.
  EXPECT_FALSE(flow_controller_->FlowControlViolation());
  EXPECT_EQ(kInitialSessionFlowControlWindowForTest,
            QuicFlowControllerPeer::ReceiveWindowSize(flow_controller_.get()));

  // Move time forward, but by more than two RTTs.  Then receive and consume
  // some more, forcing a second WINDOW_UPDATE with unchanged max window size.
  connection_.AdvanceTime(QuicTime::Delta::FromMilliseconds(2 * kRtt + 1));

  receive_offset += threshold + 1;
  EXPECT_TRUE(flow_controller_->UpdateHighestReceivedOffset(receive_offset));

  flow_controller_->AddBytesConsumed(threshold + 1);
  EXPECT_FALSE(flow_controller_->FlowControlViolation());

  QuicByteCount new_threshold =
      QuicFlowControllerPeer::WindowUpdateThreshold(flow_controller_.get());

  EXPECT_EQ(new_threshold, threshold);
}

}  // namespace test
}  // namespace net