// Copyright 2023 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "ash/frame_sink/frame_sink_holder.h"

#include <algorithm>
#include <memory>
#include <vector>

#include "ash/constants/ash_features.h"
#include "ash/frame_sink/frame_sink_holder_test_api.h"
#include "ash/frame_sink/frame_sink_host.h"
#include "ash/frame_sink/test/test_begin_frame_source.h"
#include "ash/frame_sink/test/test_frame_factory.h"
#include "ash/frame_sink/test/test_layer_tree_frame_sink.h"
#include "ash/frame_sink/ui_resource_manager.h"
#include "ash/shell.h"
#include "ash/test/ash_test_base.h"
#include "base/functional/bind.h"
#include "base/memory/ptr_util.h"
#include "base/memory/raw_ptr.h"
#include "base/memory/weak_ptr.h"
#include "base/run_loop.h"
#include "base/test/scoped_feature_list.h"
#include "cc/resources/resource_pool.h"
#include "components/viz/client/client_resource_provider.h"
#include "components/viz/common/frame_sinks/begin_frame_args.h"
#include "components/viz/common/frame_sinks/begin_frame_source.h"
#include "components/viz/common/gpu/context_provider.h"
#include "components/viz/common/quads/compositor_frame.h"
#include "components/viz/common/resources/resource_id.h"
#include "components/viz/common/resources/returned_resource.h"
#include "components/viz/common/resources/transferable_resource.h"
#include "gpu/command_buffer/client/test_shared_image_interface.h"
#include "testing/gmock/include/gmock/gmock.h"
#include "ui/gfx/geometry/rect.h"
#include "ui/gfx/geometry/size.h"

namespace ash {
namespace {

class StubBeginFrameSource : public viz::BeginFrameSource {
 public:
  StubBeginFrameSource() : viz::BeginFrameSource(0u) {}

  void DidFinishFrame(viz::BeginFrameObserver* obs,
                      viz::DisplaySchedulerDrawResult result) override {}
  void AddObserver(viz::BeginFrameObserver* obs) override {}
  void RemoveObserver(viz::BeginFrameObserver* obs) override {}
  void OnGpuNoLongerBusy() override {}
};

MATCHER_P(IsBeginFrameAckEqual, value, "") {
  return arg.frame_id == value.frame_id && arg.trace_id == value.trace_id &&
         arg.has_damage == value.has_damage;
}

class FrameSinkHolderTest : public AshTestBase,
                            public testing::WithParamInterface<bool> {
 public:
  FrameSinkHolderTest() {
    scoped_feature_list_.InitWithFeatureState(
        features::kFrameSinkHostNewBackend, GetParam());
  }
  FrameSinkHolderTest(const FrameSinkHolderTest&) = delete;
  FrameSinkHolderTest& operator=(const FrameSinkHolderTest&) = delete;

  // AshTestBase:
  void SetUp() override {
    AshTestBase::SetUp();
    aura::Window* root_window =
        Shell::Get()->GetRootWindowForDisplayId(GetPrimaryDisplay().id());
    auto host_window = std::make_unique<aura::Window>(/*delegate=*/nullptr);
    host_window_ = host_window.release();
    host_window_->Init(ui::LayerType::LAYER_SOLID_COLOR);
    root_window->AddChild(host_window_);

    frame_factory_ = std::make_unique<TestFrameFactory>();

    auto layer_tree_frame_sink = std::make_unique<TestLayerTreeFrameSink>();

    frame_sink_holder_ = std::make_unique<FrameSinkHolder>(
        std::move(layer_tree_frame_sink),
        base::BindRepeating(&TestFrameFactory::CreateCompositorFrame,
                            base::Unretained(frame_factory_.get())),
        /*on_first_frame_requested_callback=*/base::DoNothing(),
        /*on_frame_sink_lost_callback=*/base::DoNothing());

    holder_weak_ptr_ = frame_sink_holder_->GetWeakPtr();

    sii_ = base::MakeRefCounted<gpu::TestSharedImageInterface>();
  }

  void TearDown() override {
    frame_sink_holder_.reset();
    frame_factory_.reset();
    host_window_ = nullptr;
    AshTestBase::TearDown();
  }

  TestLayerTreeFrameSink* layer_tree_frame_sink() {
    return static_cast<TestLayerTreeFrameSink*>(
        holder_weak_ptr_->layer_tree_frame_sink_for_test());
  }

  size_t GetExportedResourcesCount(FrameSinkHolder* holder) {
    if (features::IsFrameSinkHostNewBackendEnabled()) {
      FrameSinkHolderTestApi test_api(holder);
      const size_t count = test_api.GetExportedResourcesCount();
      EXPECT_EQ(
          test_api.client_resource_provider()->num_resources_for_testing(),
          count);
      EXPECT_GE(test_api.resource_pool()->GetTotalResourceCountForTesting(),
                count);
      return count;
    }
    return holder->resource_manager().exported_resources_count();
  }

  viz::TransferableResource CreateAndExportResource() {
    if (features::IsFrameSinkHostNewBackendEnabled()) {
      FrameSinkHolderTestApi test_api(frame_sink_holder_.get());
      auto resource = test_api.resource_pool()->AcquireResource(
          gfx::Size(20, 20), viz::SinglePlaneFormat::kBGRA_8888,
          gfx::ColorSpace());
      resource.InstallSoftwareBacking(sii_, "TestFrame");
      test_api.resource_pool()->PrepareForExport(
          resource, viz::TransferableResource::ResourceSource::kTest);

      std::vector<viz::TransferableResource> transferable_resources;
      test_api.client_resource_provider()->PrepareSendToParent(
          {resource.resource_id_for_export()}, &transferable_resources,
          sii_.get());

      test_api.resource_pool()->ReleaseResource(std::move(resource));
      return transferable_resources[0];
    }

    return GetResourceManager().OfferAndPrepareResourceForExport(
        MakeResource());
  }

  std::unique_ptr<UiResource> MakeResource() {
    const gfx::Size kSize = gfx::Size(20, 20);
    auto shared_image = sii_->CreateSharedImage(
        {viz::SinglePlaneFormat::kBGRA_8888, kSize, gfx::ColorSpace(),
         gpu::SHARED_IMAGE_USAGE_DISPLAY_READ, "FastInkRootViewFrame"},
        gpu::kNullSurfaceHandle);

    auto resource = std::make_unique<UiResource>(sii_, std::move(shared_image));
    resource->ui_source_id = 1u;
    return resource;
  }

  UiResourceManager& GetResourceManager() {
    return frame_sink_holder_->resource_manager();
  }

  scoped_refptr<gpu::SharedImageInterface> sii_;

  // If `frame_sink_holder_` lifetime has been extended in a unittest and the
  // holder did not schedule a delete task, it will get destroyed once we
  // delete the root_window of `host_window_`.
  std::unique_ptr<FrameSinkHolder> frame_sink_holder_;
  raw_ptr<aura::Window> host_window_;

  // Will be used to access the frame_sink_holder once we pass the
  // ownership of `frame_sink_holder_` to
  // `DeleteWhenLastResourceHasBeenReclaimed()` in unittests.
  base::WeakPtr<FrameSinkHolder> holder_weak_ptr_;

  // Factory to create test compositor frames.
  std::unique_ptr<TestFrameFactory> frame_factory_;

  base::test::ScopedFeatureList scoped_feature_list_;
};

TEST_P(FrameSinkHolderTest, SubmitFrameSynchronouslyBeforeFirstFrameRequested) {
  FrameSinkHolderTestApi test_api(frame_sink_holder_.get());

  frame_factory_->SetFrameMetaData(gfx::Size(100, 100), 1.0);

  ASSERT_FALSE(test_api.IsFirstFrameRequested());
  frame_sink_holder_->SubmitCompositorFrame(/*synchronous_draw=*/true);

  // Confirm that FrameSinkHolder did not submit any frame yet.
  EXPECT_TRUE(test_api.LastSubmittedFrameSize().IsEmpty());
  EXPECT_EQ(layer_tree_frame_sink()->num_of_frames_received(), 0);

  // FrameSinkHolder has pending a frame that will be sent out asynchronously.
  EXPECT_TRUE(test_api.IsPendingFrame());

  // Asynchronous frame request.
  frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());
  ASSERT_TRUE(test_api.IsFirstFrameRequested());
  EXPECT_FALSE(test_api.IsPendingFrame());
  EXPECT_TRUE(test_api.IsPendingFrameAck());

  // LayerTreeFrameSink received the frame.
  EXPECT_EQ(layer_tree_frame_sink()->num_of_frames_received(), 1);

  // Manual BeginFrameAck is used for synchronous frames only.
  EXPECT_THAT(layer_tree_frame_sink()
                  ->GetLatestReceivedFrame()
                  .metadata.begin_frame_ack,
              testing::Not(IsBeginFrameAckEqual(
                  viz::BeginFrameAck::CreateManualAckWithDamage())));

  frame_sink_holder_->DidReceiveCompositorFrameAck();
  EXPECT_FALSE(test_api.IsPendingFrameAck());

  layer_tree_frame_sink()->ResetLatestFrameState();

  // Now that FrameSinkHolder has received the requested for the first frame, it
  // can now submit frames synchronously.
  frame_sink_holder_->SubmitCompositorFrame(/*synchronous_draw=*/true);
  EXPECT_EQ(layer_tree_frame_sink()->num_of_frames_received(), 2);
  EXPECT_EQ(test_api.LastSubmittedFrameSize(),
            layer_tree_frame_sink()->GetLatestReceivedFrame().size_in_pixels());

  // Manual BeginFrameAck is used only for synchronously submitted frames.
  EXPECT_THAT(
      layer_tree_frame_sink()
          ->GetLatestReceivedFrame()
          .metadata.begin_frame_ack,
      IsBeginFrameAckEqual(viz::BeginFrameAck::CreateManualAckWithDamage()));
}

TEST_P(FrameSinkHolderTest, ObserveBeginFrameSourceOnDemand) {
  FrameSinkHolderTestApi test_api(frame_sink_holder_.get());

  StubBeginFrameSource source;
  frame_sink_holder_->SetBeginFrameSource(&source);

  // FrameSinkHolder should be observing the source when it is set.
  EXPECT_TRUE(test_api.IsObservingBeginFrameSource());

  // After consecutively not producing frames for a certain number of
  // BeginFrames, FrameSinkHolder should stop observing the source.
  for (int i = 0; i < 4; i++) {
    frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());
  }

  frame_sink_holder_->SubmitCompositorFrame(/*synchronous_draw=*/true);
  frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());
  EXPECT_TRUE(test_api.IsObservingBeginFrameSource());

  for (int i = 0; i < 5; i++) {
    frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());
  }

  EXPECT_FALSE(test_api.IsObservingBeginFrameSource());

  // However, if there is request to submit a new frame, FrameSinkHolder should
  // start observing the source again.
  frame_sink_holder_->SubmitCompositorFrame(/*synchronous_draw=*/true);
  EXPECT_TRUE(test_api.IsObservingBeginFrameSource());

  frame_sink_holder_->SetBeginFrameSource(nullptr);
}

TEST_P(FrameSinkHolderTest, ObserveBeginFrameSourceOnDemand_AutoUpdate) {
  FrameSinkHolderTestApi test_api(frame_sink_holder_.get());

  frame_sink_holder_->SetAutoUpdateMode(true);

  StubBeginFrameSource source;
  frame_sink_holder_->SetBeginFrameSource(&source);

  // FrameSinkHolder should be observing BeginFrameSource when it is set.
  EXPECT_TRUE(test_api.IsObservingBeginFrameSource());

  // When auto update mode is on, we should not stop observation of
  // BeginFrameSource.
  for (int i = 0; i < 10; i++) {
    frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());
  }

  EXPECT_TRUE(test_api.IsObservingBeginFrameSource());

  // However once the auto update mode is turned off, we should stop observing
  // the BeginFrameSource as needed.
  frame_sink_holder_->SetAutoUpdateMode(false);

  for (int i = 0; i < 5; i++) {
    frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());
  }

  EXPECT_FALSE(test_api.IsObservingBeginFrameSource());

  // Renablending the auto update mode should start the observation of
  // BeginFrameSource.
  frame_sink_holder_->SetAutoUpdateMode(true);
  EXPECT_TRUE(test_api.IsObservingBeginFrameSource());

  frame_sink_holder_->SetBeginFrameSource(nullptr);
}

TEST_P(FrameSinkHolderTest, SubmitFrameSynchronouslyWhilePendingFrameAck) {
  FrameSinkHolderTestApi test_api(frame_sink_holder_.get());

  frame_factory_->SetFrameMetaData(gfx::Size(100, 100), 1.0);

  // Call OnBeginFrame so that FrameSinkHolder can know that it can submit
  // frames synchronously.
  frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());

  frame_sink_holder_->SubmitCompositorFrame(/*synchronous_draw=*/true);
  EXPECT_TRUE(test_api.IsPendingFrameAck());
  EXPECT_EQ(layer_tree_frame_sink()->num_of_frames_received(), 1);

  frame_factory_->SetFrameMetaData(gfx::Size(200, 200), 1.0);
  frame_sink_holder_->SubmitCompositorFrame(/*synchronous_draw=*/true);

  // This confirms that FrameSinkHolder did not submit frame synchronously,
  // since it has not received frame ack for the last frame.
  EXPECT_EQ(layer_tree_frame_sink()->GetLatestReceivedFrame().size_in_pixels(),
            gfx::Size(100, 100));
  EXPECT_EQ(layer_tree_frame_sink()->num_of_frames_received(), 1);

  // FrameSinkHolder fell to asynchronous frame submission.
  EXPECT_TRUE(test_api.IsPendingFrame());
}

TEST_P(FrameSinkHolderTest, HandlingAsynchronousFrameRequests_NoAutoUpdate) {
  FrameSinkHolderTestApi test_api(frame_sink_holder_.get());

  frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());

  // FrameSinkHolder has no request to submit a frame.
  auto skipped_reason = layer_tree_frame_sink()->GetLatestFrameSkippedReason();
  ASSERT_TRUE(skipped_reason.has_value());
  EXPECT_EQ(skipped_reason, cc::FrameSkippedReason::kNoDamage);

  frame_factory_->SetFrameMetaData(gfx::Size(100, 100), 1.0);

  frame_sink_holder_->SubmitCompositorFrame(/*synchronous_draw=*/false);
  EXPECT_TRUE(test_api.IsPendingFrame());

  // This time FrameSinkHolder has a request to submit frame asynchronously.
  frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());
  EXPECT_EQ(layer_tree_frame_sink()->num_of_frames_received(), 1);

  // Asynchronously submitted frames will not have manual BeginFrameAck.
  EXPECT_THAT(layer_tree_frame_sink()
                  ->GetLatestReceivedFrame()
                  .metadata.begin_frame_ack,
              testing::Not(IsBeginFrameAckEqual(
                  viz::BeginFrameAck::CreateManualAckWithDamage())));
  EXPECT_FALSE(test_api.IsPendingFrame());
  EXPECT_TRUE(test_api.IsPendingFrameAck());

  layer_tree_frame_sink()->ResetLatestFrameState();

  // FrameSinkHolder did not submit a frame since it is still waiting for ack.
  frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());

  skipped_reason = layer_tree_frame_sink()->GetLatestFrameSkippedReason();
  ASSERT_TRUE(skipped_reason.has_value());
  EXPECT_EQ(skipped_reason, cc::FrameSkippedReason::kWaitingOnMain);

  EXPECT_EQ(layer_tree_frame_sink()->num_of_frames_received(), 1);

  // Received ack.
  frame_sink_holder_->DidReceiveCompositorFrameAck();
  EXPECT_FALSE(test_api.IsPendingFrameAck());

  layer_tree_frame_sink()->ResetLatestFrameState();

  // FrameSinkHolder did not submit anything because it did not have any pending
  // request.
  frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());

  skipped_reason = layer_tree_frame_sink()->GetLatestFrameSkippedReason();
  ASSERT_TRUE(skipped_reason.has_value());
  EXPECT_EQ(skipped_reason, cc::FrameSkippedReason::kNoDamage);

  EXPECT_EQ(layer_tree_frame_sink()->num_of_frames_received(), 1);

  // FrameSinkHolder has an async request again.
  frame_sink_holder_->SubmitCompositorFrame(/*synchronous_draw=*/false);
  EXPECT_TRUE(test_api.IsPendingFrame());
  EXPECT_FALSE(test_api.IsPendingFrameAck());

  layer_tree_frame_sink()->ResetLatestFrameState();

  // FrameSinkHolder should now submit a new frame.
  frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());

  // Asynchronously submitted frames will not have manual begin_frame_ack.
  EXPECT_THAT(layer_tree_frame_sink()
                  ->GetLatestReceivedFrame()
                  .metadata.begin_frame_ack,
              testing::Not(IsBeginFrameAckEqual(
                  viz::BeginFrameAck::CreateManualAckWithDamage())));
  EXPECT_EQ(layer_tree_frame_sink()->num_of_frames_received(), 2);
  EXPECT_FALSE(test_api.IsPendingFrame());
  EXPECT_TRUE(test_api.IsPendingFrameAck());
}

TEST_P(FrameSinkHolderTest, DontSubmitNewFramesWhenWaitingToDeleteSinkHolder) {
  FrameSinkHolderTestApi test_api(frame_sink_holder_.get());
  base::RunLoop loop;

  auto resource_1 = CreateAndExportResource();

  frame_factory_->SetFrameResources({resource_1});
  frame_factory_->SetFrameMetaData(gfx::Size(100, 100), 1.0);

  // Call OnBeginFrame so that FrameSinkHolder can know that it can submit
  // frames synchronously.
  frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());
  frame_sink_holder_->SubmitCompositorFrame(/*synchronous_draw=*/true);
  EXPECT_EQ(layer_tree_frame_sink()->num_of_frames_received(), 1);

  // The lifetime of frame_sink_holder has been extended since there are still
  // some exported resources.
  EXPECT_FALSE(FrameSinkHolder::DeleteWhenLastResourceHasBeenReclaimed(
      std::move(frame_sink_holder_), host_window_));

  ASSERT_TRUE(holder_weak_ptr_);

  // During deletion FrameSinkHolder submits a empty frame.
  EXPECT_EQ(layer_tree_frame_sink()->num_of_frames_received(), 2);

  layer_tree_frame_sink()->ResetLatestFrameState();

  holder_weak_ptr_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());

  // Confirms that FrameSinkHolder did not submit a new frame on asynchronous
  // request.
  EXPECT_EQ(layer_tree_frame_sink()->num_of_frames_received(), 2);
}

TEST_P(FrameSinkHolderTest,
       DeleteSinkHolderImmediatelyWhenNoFramesIsSubmitted) {
  FrameSinkHolderTestApi test_api(frame_sink_holder_.get());

  // Confirms that FrameSinkHolder has not submitted any frames yet.
  EXPECT_TRUE(test_api.LastSubmittedFrameSize().IsEmpty());
  EXPECT_EQ(layer_tree_frame_sink()->num_of_frames_received(), 0);

  // FrameSinkHolder will get deleted straight away since it has not submitted
  // any resources to the display compositor.
  EXPECT_TRUE(FrameSinkHolder::DeleteWhenLastResourceHasBeenReclaimed(
      std::move(frame_sink_holder_), host_window_));

  // Since FrameSinkHolder is deleted immediately, we expect the weak_ptr to be
  // not valid.
  EXPECT_FALSE(holder_weak_ptr_);
}

TEST_P(FrameSinkHolderTest, ExtendLifeTimeOfHolderToRootWindow) {
  FrameSinkHolderTestApi test_api(frame_sink_holder_.get());

  auto resource_1 = CreateAndExportResource();
  auto resource_2 = CreateAndExportResource();
  auto resource_3 = CreateAndExportResource();

  frame_factory_->SetFrameResources({resource_1, resource_2, resource_3});
  frame_factory_->SetFrameMetaData(gfx::Size(100, 100), 1.0);

  // Call OnBeginFrame so that FrameSinkHolder can know that it can submit
  // frames synchronously.
  frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());
  frame_sink_holder_->SubmitCompositorFrame(/*synchronous_draw=*/true);

  // Confirms that FrameSinkHolder has not submitted any frames.
  EXPECT_FALSE(test_api.LastSubmittedFrameSize().IsEmpty());
  EXPECT_EQ(layer_tree_frame_sink()->num_of_frames_received(), 1);

  // Since FrameSinkHolder has not received the resources back from display
  // compositor, it extend its lifetime.
  EXPECT_FALSE(FrameSinkHolder::DeleteWhenLastResourceHasBeenReclaimed(
      std::move(frame_sink_holder_), host_window_));

  // Since FrameSinkHolder lifetime is extend, we expect the weak_ptr to be
  // valid.
  EXPECT_TRUE(holder_weak_ptr_);
}

TEST_P(FrameSinkHolderTest, KeepSubmittingFrameWhenAutoUpdateIsOn) {
  FrameSinkHolderTestApi test_api(frame_sink_holder_.get());

  frame_factory_->SetFrameMetaData(gfx::Size(100, 100), 1.0);

  // Request a frame.
  frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());

  // Since auto_fresh_mode is off, FrameSinkHolder did not submit any frame as
  // there was not request for a frame submission.
  EXPECT_EQ(layer_tree_frame_sink()->num_of_frames_received(), 0);

  // Request a frame again. FrameSinkHolder should not submit a frame.
  frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());
  EXPECT_EQ(layer_tree_frame_sink()->num_of_frames_received(), 0);

  frame_sink_holder_->SetAutoUpdateMode(/*mode=*/true);

  // After auto_fresh_mode on, when compositor requests for a frame,
  // FrameSinkHolder should submit a frame.
  frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());
  EXPECT_EQ(layer_tree_frame_sink()->num_of_frames_received(), 1);

  // FrameSinkHolder should not submit a new frame sas it has no received an
  // ack from the compositor,
  frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());
  EXPECT_EQ(layer_tree_frame_sink()->num_of_frames_received(), 1);

  // Receive an ack.
  frame_sink_holder_->DidReceiveCompositorFrameAck();

  // In auto_fresh mode, FrameSinkHolder will keep on submitting frames
  // asynchronously.
  frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());
  EXPECT_EQ(layer_tree_frame_sink()->num_of_frames_received(), 2);
  frame_sink_holder_->DidReceiveCompositorFrameAck();
  frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());
  EXPECT_EQ(layer_tree_frame_sink()->num_of_frames_received(), 3);
}

TEST_P(FrameSinkHolderTest, DeleteHolderAfterReclaimingAllResources) {
  FrameSinkHolderTestApi test_api(frame_sink_holder_.get());
  base::RunLoop loop;

  auto resource_1 = CreateAndExportResource();
  auto resource_2 = CreateAndExportResource();

  frame_factory_->SetFrameResources({resource_1, resource_2});
  frame_factory_->SetFrameMetaData(gfx::Size(100, 100), 1.0);

  // Call OnBeginFrame so that FrameSinkHolder can know that it can submit
  // frames synchronously.
  frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());
  frame_sink_holder_->SubmitCompositorFrame(/*synchronous_draw=*/true);

  EXPECT_FALSE(FrameSinkHolder::DeleteWhenLastResourceHasBeenReclaimed(
      std::move(frame_sink_holder_), host_window_));

  // The lifetime of frame_sink_holder has been extended since there are still
  // some exported resources.
  ASSERT_TRUE(holder_weak_ptr_);

  std::vector<viz::ReturnedResource> to_be_returned_resources;
  layer_tree_frame_sink()->GetFrameResourcesToReturn(to_be_returned_resources);

  // Reclaim the exported resources.
  holder_weak_ptr_->ReclaimResources(std::move(to_be_returned_resources));

  // Wait for the deletion task to be completed.
  loop.RunUntilIdle();

  ASSERT_FALSE(holder_weak_ptr_);
}

TEST_P(FrameSinkHolderTest, LayerTreeFrameSinkLost) {
  FrameSinkHolderTestApi test_api(frame_sink_holder_.get());

  auto resource_1 = CreateAndExportResource();

  frame_factory_->SetFrameResources({resource_1});
  frame_factory_->SetFrameMetaData(gfx::Size(100, 100), 1.0);

  // Call OnBeginFrame so that FrameSinkHolder can know that it can submit
  // frames synchronously.
  frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());
  frame_sink_holder_->SubmitCompositorFrame(/*synchronous_draw=*/true);

  EXPECT_EQ(GetExportedResourcesCount(frame_sink_holder_.get()), 1u);

  frame_sink_holder_->DidLoseLayerTreeFrameSink();

  EXPECT_EQ(GetExportedResourcesCount(frame_sink_holder_.get()), 0u);
}

TEST_P(FrameSinkHolderTest,
       LayerTreeFrameSinkLostWhenWaitingToDeleteSinkHolder) {
  FrameSinkHolderTestApi test_api(frame_sink_holder_.get());
  base::RunLoop loop;

  auto resource_1 = CreateAndExportResource();

  frame_factory_->SetFrameResources({resource_1});
  frame_factory_->SetFrameMetaData(gfx::Size(100, 100), 1.0);

  // Call OnBeginFrame so that FrameSinkHolder can know that it can submit
  // frames synchronously.
  frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());
  frame_sink_holder_->SubmitCompositorFrame(/*synchronous_draw=*/true);

  EXPECT_FALSE(FrameSinkHolder::DeleteWhenLastResourceHasBeenReclaimed(
      std::move(frame_sink_holder_), host_window_));

  // The lifetime of frame_sink_holder has been extended since there are still
  // some exported resources.
  ASSERT_TRUE(holder_weak_ptr_);

  holder_weak_ptr_->DidLoseLayerTreeFrameSink();

  // Since FrameSinkHolder cannot reclaim back exported resources, it schedules
  // to delete itself.
  // Wait for deletion task to complete.
  loop.RunUntilIdle();

  ASSERT_FALSE(holder_weak_ptr_);
}

TEST_P(FrameSinkHolderTest,
       DeleteSinkHolderWithExportedResources_DuringShutdown) {
  FrameSinkHolderTestApi test_api(frame_sink_holder_.get());

  auto resource_1 = CreateAndExportResource();

  frame_factory_->SetFrameResources({resource_1});
  frame_factory_->SetFrameMetaData(gfx::Size(100, 100), 1.0);

  // Call OnBeginFrame so that FrameSinkHolder can know that it can submit
  // frames synchronously.
  frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());
  frame_sink_holder_->SubmitCompositorFrame(/*synchronous_draw=*/true);

  // Confirms we have an exported resource.
  EXPECT_EQ(GetExportedResourcesCount(frame_sink_holder_.get()), 1u);

  // During shutdown, root_window can be null. We can replicate it by
  // removing the host window from the window hierarchy.
  aura::Window* root_window =
      Shell::Get()->GetRootWindowForDisplayId(GetPrimaryDisplay().id());
  root_window->RemoveChild(host_window_);

  // Since `host_window_` is removed from window tree hierarchy, wrapping it in
  // a unique_ptr to delete this object as it goes out of scope and stop it
  // from leaking memory.
  auto host_window = base::WrapUnique<aura::Window>(host_window_);
  host_window_ = nullptr;

  // Since FrameSinkHolder cannot extend its lifetime, it marks the resources
  // as lost and deletes itself immediately.
  EXPECT_TRUE(FrameSinkHolder::DeleteWhenLastResourceHasBeenReclaimed(
      std::move(frame_sink_holder_), host_window.get()));

  // Since FrameSinkHolder is deleted immediately, we expect the weak_ptr to be
  // not valid.
  EXPECT_FALSE(holder_weak_ptr_);
}

TEST_P(FrameSinkHolderTest,
       DeleteSinkHolderImmediatelyWhenNoExportedResources) {
  FrameSinkHolderTestApi test_api(frame_sink_holder_.get());

  auto resource_1 = CreateAndExportResource();

  frame_factory_->SetFrameResources({resource_1});
  frame_factory_->SetFrameMetaData(gfx::Size(100, 100), 1.0);

  // Call OnBeginFrame so that FrameSinkHolder can know that it can submit
  // frames synchronously.
  frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());
  frame_sink_holder_->SubmitCompositorFrame(/*synchronous_draw=*/true);

  // Confirms that FrameSinkHolder has submitted a frame.
  EXPECT_FALSE(test_api.LastSubmittedFrameSize().IsEmpty());
  EXPECT_EQ(layer_tree_frame_sink()->num_of_frames_received(), 1);

  std::vector<viz::ReturnedResource> to_be_returned_resources;
  layer_tree_frame_sink()->GetFrameResourcesToReturn(to_be_returned_resources);

  frame_sink_holder_->ReclaimResources(std::move(to_be_returned_resources));

  // We can delete the holder straight way since we have no exported resources.
  ASSERT_EQ(GetExportedResourcesCount(frame_sink_holder_.get()), 0u);

  EXPECT_TRUE(FrameSinkHolder::DeleteWhenLastResourceHasBeenReclaimed(
      std::move(frame_sink_holder_), host_window_));

  // Since FrameSinkHolder is deleted immediately, we expect the weak_ptr to be
  // not valid.
  EXPECT_FALSE(holder_weak_ptr_);
}

TEST_P(FrameSinkHolderTest, DeleteSinkHolderImmediatelyWhenFrameSinkIsLost) {
  FrameSinkHolderTestApi test_api(frame_sink_holder_.get());

  auto resource_1 = CreateAndExportResource();

  frame_factory_->SetFrameResources({resource_1});
  frame_factory_->SetFrameMetaData(gfx::Size(100, 100), 1.0);

  // Call OnBeginFrame so that FrameSinkHolder can know that it can submit
  // frames synchronously.
  frame_sink_holder_->OnBeginFrame(CreateValidBeginFrameArgsForTesting());
  frame_sink_holder_->SubmitCompositorFrame(/*synchronous_draw=*/true);

  // Confirms that FrameSinkHolder has submitted a frame.
  EXPECT_FALSE(test_api.LastSubmittedFrameSize().IsEmpty());
  EXPECT_EQ(layer_tree_frame_sink()->num_of_frames_received(), 1);

  ASSERT_EQ(GetExportedResourcesCount(frame_sink_holder_.get()), 1u);

  frame_sink_holder_->DidLoseLayerTreeFrameSink();

  // Since frame sink is lost, the exported resources cannot be reclaimed.
  EXPECT_TRUE(FrameSinkHolder::DeleteWhenLastResourceHasBeenReclaimed(
      std::move(frame_sink_holder_), host_window_));

  // Since FrameSinkHolder is deleted immediately, we expect the weak_ptr to be
  // not valid.
  EXPECT_FALSE(holder_weak_ptr_);
}

INSTANTIATE_TEST_SUITE_P(All, FrameSinkHolderTest, testing::Bool());

}  // namespace
}  // namespace ash