// Copyright 2022 The Fuchsia 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 <fidl/fuchsia.kernel/cpp/fidl.h>
#include <lib/component/incoming/cpp/protocol.h>

#include "lib/zx/time.h"
#include "src/lib/fxl/strings/string_printf.h"
#include "src/virtualization/tests/lib/guest_test.h"

namespace {
// 5.5.6.4 Memory Statistics Tags
constexpr uint16_t VIRTIO_BALLOON_S_MEMFREE = 4;
constexpr uint16_t VIRTIO_BALLOON_S_MEMTOT = 5;
constexpr uint16_t VIRTIO_BALLOON_S_AVAIL = 6;

constexpr size_t kVirtioBalloonPageCount = 256;

template <class T>
class BalloonGuestTest : public GuestTest<T> {
 public:
  BalloonGuestTest() {
    auto client_end = component::Connect<fuchsia_kernel::Stats>();
    if (!client_end.is_ok()) {
      FX_PLOGS(ERROR, client_end.status_value()) << "Failed to connect to kernel stats";
    }
    stats.Bind(std::move(*client_end));
  }

  std::unordered_map<uint16_t, uint64_t> GetMemStats(
      const fuchsia::virtualization::BalloonControllerSyncPtr& balloon_controller) {
    ::fidl::VectorPtr<::fuchsia::virtualization::MemStat> mem_stats;
    int32_t mem_stats_status = 0;
    zx_status_t status = balloon_controller->GetMemStats(&mem_stats_status, &mem_stats);
    EXPECT_EQ(status, ZX_OK);
    EXPECT_EQ(mem_stats_status, ZX_OK);
    std::unordered_map<uint16_t, uint64_t> stats;
    for (auto& el : mem_stats.value()) {
      stats[el.tag] = el.val;
    }
    return stats;
  }

  void TestGetMemStats(
      const char* trace_context,
      const fuchsia::virtualization::BalloonControllerSyncPtr& balloon_controller) {
    SCOPED_TRACE(trace_context);

    auto stats = GetMemStats(balloon_controller);
    EXPECT_GT(stats[VIRTIO_BALLOON_S_MEMTOT], 0u);
    EXPECT_GT(stats[VIRTIO_BALLOON_S_MEMFREE], 0u);
    EXPECT_GT(stats[VIRTIO_BALLOON_S_AVAIL], 0u);
    EXPECT_LE(stats[VIRTIO_BALLOON_S_MEMFREE], stats[VIRTIO_BALLOON_S_MEMTOT]);
    EXPECT_LE(stats[VIRTIO_BALLOON_S_AVAIL], stats[VIRTIO_BALLOON_S_MEMTOT]);
  }

  void PrintKernelMemoryStats(std::string_view header) {
    auto memory_stats = stats->GetMemoryStats();
    FX_LOGS(INFO) << header
                  << ": total_bytes=" << memory_stats->stats().total_bytes().value() / 1024 / 1024
                  << " MiB free_bytes=" << memory_stats->stats().free_bytes().value() / 1024 / 1024
                  << " MiB";
  }

  fidl::SyncClient<fuchsia_kernel::Stats> stats;
};

// Zircon does not yet have a virtio balloon driver.
using GuestTypes = ::testing::Types<DebianEnclosedGuest>;
TYPED_TEST_SUITE(BalloonGuestTest, GuestTypes, GuestTestNameGenerator);

TYPED_TEST(BalloonGuestTest, InflateDeflate) {
  std::string result;
  EXPECT_EQ(this->Execute({"echo", "test"}, &result), ZX_OK);
  EXPECT_EQ(result, "test\n");

  fuchsia::virtualization::BalloonControllerSyncPtr balloon_controller;
  ASSERT_TRUE(this->ConnectToBalloon(balloon_controller.NewRequest()));

  uint32_t initial_num_pages;
  uint32_t requested_num_pages;
  zx_status_t status = balloon_controller->GetBalloonSize(&initial_num_pages, &requested_num_pages);
  ASSERT_EQ(status, ZX_OK);
  EXPECT_EQ(requested_num_pages, initial_num_pages);
  this->TestGetMemStats("Before inflate", balloon_controller);

  // Request an increase to the number of pages in the balloon.
  status = balloon_controller->RequestNumPages(initial_num_pages + kVirtioBalloonPageCount);
  ASSERT_EQ(status, ZX_OK);

  // Verify that the number of pages eventually equals the requested number. The
  // guest may not respond to the request immediately so we call GetBalloonSize in
  // a loop.
  uint32_t current_num_pages;
  while (true) {
    status = balloon_controller->GetBalloonSize(&current_num_pages, &requested_num_pages);
    ASSERT_EQ(status, ZX_OK);
    EXPECT_EQ(requested_num_pages, initial_num_pages + kVirtioBalloonPageCount);
    if (current_num_pages == initial_num_pages + kVirtioBalloonPageCount) {
      break;
    }
  }
  this->TestGetMemStats("After inflate", balloon_controller);

  // Request a decrease to the number of pages in the balloon back to the
  // initial value.
  status = balloon_controller->RequestNumPages(initial_num_pages);
  ASSERT_EQ(status, ZX_OK);

  while (true) {
    status = balloon_controller->GetBalloonSize(&current_num_pages, &requested_num_pages);
    ASSERT_EQ(status, ZX_OK);
    EXPECT_EQ(requested_num_pages, initial_num_pages);
    if (current_num_pages == initial_num_pages) {
      break;
    }
  }
  this->TestGetMemStats("After deflate", balloon_controller);
}

TYPED_TEST(BalloonGuestTest, VirtioBalloonFreePageReporting) {
  const uint64_t starting_free_memory_mib =
      this->stats->GetMemoryStats()->stats().free_bytes().value() / 1024 / 1024;

  // Allocate 256MiB or 50% of the free memory, whatever is the smallest.
  // We don't want to cause the memory pressure warning by allocating too
  // much memory
  const uint64_t max_alloc_amount_mib = 256;
  const uint64_t alloc_amount_mib =
      std::min(max_alloc_amount_mib, starting_free_memory_mib / 10 * 5);
  this->PrintKernelMemoryStats("Before the guest alloc");

  FX_LOGS(INFO) << fxl::StringPrintf("Allocate and release %lu MiB in the guest", alloc_amount_mib);
  // This call will allocate and immediate release the specified amount of
  // memory in the guest.
  // From the guest perspective, memory is available once it got released.
  // From the host perspective, memory is taken by the guest and not available
  // until it will be reclaimed by the free page reporting.
  std::string result;
  ASSERT_EQ(this->RunUtil(
                fxl::StringPrintf("memory_test_util alloc --size-mb 1 --num %lu", alloc_amount_mib),
                {}, &result),
            ZX_OK);

  this->PrintKernelMemoryStats("After the guest alloc and release");

  // Require 10% of allocated memory to be reclaimed to detect free page reporting
  // Requiring 50% was causing occasional flakes, especially when memory was low to begin with.
  // TODO(https://fxbug.dev/42063850) Remove added logging during the reclaim wait once flake is resolved
  const uint64_t reclaim_success_threshold =
      starting_free_memory_mib - alloc_amount_mib + alloc_amount_mib / 10;
  FX_LOGS(INFO) << "Waiting for the virtio balloon to reclaim memory. reclaim_success_threshold="
                << reclaim_success_threshold;
  const zx::time deadline = zx::deadline_after(zx::sec(30));

  while (zx::clock::get_monotonic() < deadline &&
         reclaim_success_threshold >
             this->stats->GetMemoryStats()->stats().free_bytes().value() / 1024 / 1024) {
    zx::nanosleep(zx::deadline_after(zx::msec(100)));
    this->PrintKernelMemoryStats("Waiting for memory reclaim");
  }
  this->PrintKernelMemoryStats("After the memory reclaim");
  // Prefer explicit fail instead of getting stuck in the loop above
  // Fail the test with extra logging if host's free memory didn't get above the threshold
  ZX_ASSERT(zx::clock::get_monotonic() < deadline);
}

TYPED_TEST(BalloonGuestTest, VirtioBalloonInflateOnMemoryPressureWarningDeflateOnNormal) {
  fuchsia::virtualization::BalloonControllerSyncPtr balloon_controller;
  ASSERT_TRUE(this->ConnectToBalloon(balloon_controller.NewRequest()));

  uint32_t initial_num_pages;
  uint32_t requested_num_pages;
  ASSERT_EQ(balloon_controller->GetBalloonSize(&initial_num_pages, &requested_num_pages), ZX_OK);
  EXPECT_EQ(initial_num_pages, 0u);
  EXPECT_EQ(requested_num_pages, 0u);

  FX_LOGS(INFO) << "Memory pressure warning";
  this->GetEnclosedGuest().GetMemoryPressureProvider()->OnLevelChanged(
      ::fuchsia::memorypressure::Level::WARNING);

  while (true) {
    uint32_t cur_num_pages;
    ASSERT_EQ(balloon_controller->GetBalloonSize(&cur_num_pages, &requested_num_pages), ZX_OK);

    auto mem_stats = this->GetMemStats(balloon_controller);

    if (cur_num_pages > 0u && requested_num_pages > 0u) {
      break;
    }
    zx::nanosleep(zx::deadline_after(zx::msec(100)));
  };

  FX_LOGS(INFO) << "Back to memory pressure normal";
  this->GetEnclosedGuest().GetMemoryPressureProvider()->OnLevelChanged(
      ::fuchsia::memorypressure::Level::NORMAL);

  while (true) {
    uint32_t cur_num_pages;
    ASSERT_EQ(balloon_controller->GetBalloonSize(&cur_num_pages, &requested_num_pages), ZX_OK);

    if (cur_num_pages == 0u && requested_num_pages == 0u) {
      break;
    }
    zx::nanosleep(zx::deadline_after(zx::msec(100)));
  };
}

}  // namespace