src/graphics/bin/agis/agis.cc

// Copyright 2021 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 <fuchsia/gpu/agis/cpp/fidl.h>
#include <lib/async-loop/cpp/loop.h>
#include <lib/async-loop/default.h>
#include <lib/fidl/cpp/binding_set.h>
#include <lib/sys/cpp/component_context.h>
#include <lib/syslog/cpp/macros.h>
#include <lib/zx/object.h>

#include <map>
#include <unordered_map>
#include <unordered_set>

#include <sdk/lib/syslog/cpp/log_settings.h>

namespace {
// ID provided by the client to make a vtc registration unique per client.
using ClientId = uint64_t;

// ID generated by this AGIS fidl service for each vtc registration.
using GlobalId = uint32_t;

// Value struct for |registry| map.
struct Entry {
  Entry(GlobalId global_id_in, std::string process_name_in, zx_koid_t process_koid_in,
        zx::socket vulkan_socket_in,
        fuchsia::gpu::agis::ComponentRegistry::GetVulkanSocketCallback vulkan_socket_callback_in =
            nullptr)
      : global_id(global_id_in),
        process_name(std::move(process_name_in)),
        process_koid(process_koid_in),
        vulkan_socket(vulkan_socket_in.release()),
        vulkan_socket_callback(std::move(vulkan_socket_callback_in)) {}
  GlobalId global_id;
  std::string process_name;
  zx_koid_t process_koid;
  zx::socket vulkan_socket;
  fuchsia::gpu::agis::ComponentRegistry::GetVulkanSocketCallback vulkan_socket_callback;
};

// The global registry that maps global ids to |Entry|s. Valid global ids are greater
// than zero.
std::unordered_map<GlobalId, std::shared_ptr<Entry>> registry;

// Erase a registered |id| from |registry|.
bool EraseEntry(ClientId id, std::unordered_map<ClientId, std::shared_ptr<Entry>> *session) {
  // Find |id| in the session.
  auto session_iter = session->find(id);
  if (session_iter == session->end())
    return false;
  auto global_id = session_iter->second->global_id;
  session->erase(session_iter);

  // Remove |global_id| from the |registry|.
  auto registry_iter = registry.find(global_id);
  if (registry_iter == registry.end())
    return false;
  registry.erase(registry_iter);

  return true;
}

bool PeerEndpointIsOpen(const zx::socket &endpoint) {
  if (!endpoint.is_valid()) {
    return false;
  }
  auto status =
      zx_object_wait_one(endpoint.get(), ZX_SOCKET_PEER_CLOSED, ZX_TIME_INFINITE_PAST, nullptr);

  if (status == ZX_ERR_TIMED_OUT) {
    return true;
  }

  // Report abnormal status for remaining possible return values.
  if (!(status == ZX_OK || status == ZX_ERR_CANCELED)) {
    FX_SLOG(ERROR, "PeerEndpointIsOpen", FX_KV("status", status));
  }

  return false;
}

}  // namespace

//
// Register and unregister Vulkan traceable components.
//
class ComponentRegistryImpl final : public fuchsia::gpu::agis::ComponentRegistry {
 public:
  ~ComponentRegistryImpl() override {
    for (auto session_pair : session_) {
      auto global_id = session_pair.second->global_id;
      auto i = registry.find(global_id);
      if (i == registry.end()) {
        FX_SLOG(ERROR, "Corrupt registry");
      }
      registry.erase(global_id);
    }
  }

  // Add entries to the |registry| map.
  void Register(ClientId id, zx_koid_t process_koid, std::string process_name,
                RegisterCallback callback) override {
    fuchsia::gpu::agis::ComponentRegistry_Register_Result result;

    // Test if the vtc map is full.
    if (registry.size() == fuchsia::gpu::agis::MAX_VTCS) {
      result.set_err(fuchsia::gpu::agis::Error::VTCS_EXCEEDED);
      callback(std::move(result));
      return;
    }

    // Verify that the |id| is unique.
    {
      auto i = session_.find(id);
      if (i != session_.end()) {
        result.set_err(fuchsia::gpu::agis::Error::ALREADY_REGISTERED);
        callback(std::move(result));
        return;
      }
    }

    GlobalId global_id = rand();
    while (registry.find(global_id) != registry.end() || global_id == 0) {
      global_id = rand();
    }

    // Update registry database.
    {
      auto e = std::make_shared<Entry>(global_id, process_name, process_koid,
                                       zx::socket() /* vulkan_socket */);
      registry.insert(std::make_pair(global_id, e));
      session_.insert(std::make_pair(id, e));
    }

    fuchsia::gpu::agis::ComponentRegistry_Register_Response response;
    result.set_response(response);
    callback(std::move(result));
  }

  void Unregister(ClientId id, UnregisterCallback callback) override {
    fuchsia::gpu::agis::ComponentRegistry_Unregister_Result result;

    if (!EraseEntry(id, &session_)) {
      result.set_err(fuchsia::gpu::agis::Error::NOT_FOUND);
      callback(std::move(result));
      return;
    }
    result.set_response(fuchsia::gpu::agis::ComponentRegistry_Unregister_Response());
    callback(std::move(result));
  }

  void GetVulkanSocket(ClientId id, GetVulkanSocketCallback callback) override {
    fuchsia::gpu::agis::ComponentRegistry_GetVulkanSocket_Result result;

    auto session_iter = session_.find(id);
    if (session_iter == session_.end()) {
      result.set_err(fuchsia::gpu::agis::Error::NOT_FOUND);
      callback(std::move(result));
    }

    auto &entry = session_iter->second;
    if (entry->vulkan_socket_callback) {
      result.set_err(fuchsia::gpu::agis::Error::ALREADY_REGISTERED);
      callback(std::move(result));
      return;
    }

    if (!entry->vulkan_socket.is_valid()) {
      // Initializing this vulkan_socket_callback field means the callback is pending and will
      // be satisfied when Connector::GetSocket() is called.  This is the hanging get.
      entry->vulkan_socket_callback = std::move(callback);
      return;
    }

    // |vulkan_socket| handle is valid, verify that the other endpoint is open.
    if (!PeerEndpointIsOpen(entry->vulkan_socket)) {
      // No longer open so re-register our callback.  It will now be
      // pending and is satisfied when Connector::() is called.
      entry->vulkan_socket_callback = std::move(callback);
      return;
    }

    // We were able to respond with a valid |vulkan_socket| so set the response and clear the
    // callback state.
    fuchsia::gpu::agis::ComponentRegistry_GetVulkanSocket_Response response(
        std::move(entry->vulkan_socket));
    entry->vulkan_socket_callback = nullptr;
    result.set_response(std::move(response));

    callback(std::move(result));
  }

  void AddBinding(std::unique_ptr<ComponentRegistryImpl> session,
                  fidl::InterfaceRequest<fuchsia::gpu::agis::ComponentRegistry> &&request) {
    bindings_.AddBinding(std::move(session), std::move(request));
  }

 private:
  fidl::BindingSet<fuchsia::gpu::agis::ComponentRegistry,
                   std::unique_ptr<fuchsia::gpu::agis::ComponentRegistry>>
      bindings_;

  std::unordered_map<ClientId, std::shared_ptr<Entry>> session_;
};

//
// Retrieve the list of available, registered Vulkan traceable components.
//
class ObserverImpl final : public fuchsia::gpu::agis::Observer {
 public:
  void Vtcs(VtcsCallback callback) override {
    fuchsia::gpu::agis::Observer_Vtcs_Result result;
    std::vector<fuchsia::gpu::agis::Vtc> vtcs;
    for (const auto &pair : registry) {
      auto vtc = ::fuchsia::gpu::agis::Vtc::New();
      vtc->set_global_id(pair.first);
      vtc->set_process_koid(pair.second->process_koid);
      vtc->set_process_name(pair.second->process_name);
      vtcs.emplace_back(std::move(*vtc));
    }
    fuchsia::gpu::agis::Observer_Vtcs_Response response(std::move(vtcs));
    result.set_response(std::move(response));
    callback(std::move(result));
  }

  void AddBinding(std::unique_ptr<ObserverImpl> observer,
                  fidl::InterfaceRequest<fuchsia::gpu::agis::Observer> &&request) {
    bindings_.AddBinding(std::move(observer), std::move(request));
  }

 private:
  fidl::BindingSet<fuchsia::gpu::agis::Observer, std::unique_ptr<fuchsia::gpu::agis::Observer>>
      bindings_;
};

//
// Create a Zircon socket to establish connectivity to a Vulkan traceable component.
//
class ConnectorImpl final : public fuchsia::gpu::agis::Connector {
 public:
  void GetSocket(GlobalId global_id, GetSocketCallback callback) override {
    fuchsia::gpu::agis::Connector_GetSocket_Result result;

    auto registry_iter = registry.find(global_id);
    if (registry_iter == registry.end()) {
      result.set_err(fuchsia::gpu::agis::Error::NOT_FOUND);
      callback(std::move(result));
      return;
    }

    auto &entry = registry_iter->second;
    if (entry->vulkan_socket.is_valid()) {
      result.set_err(fuchsia::gpu::agis::Error::ALREADY_REGISTERED);
      callback(std::move(result));
      return;
    }

    // Create both socket endpoints if they haven't been initialized yet.
    zx::socket ffx_socket;
    auto status = zx::socket::create(0u, &entry->vulkan_socket, &ffx_socket);
    if (status != ZX_OK) {
      FX_SLOG(ERROR, "zx::socket::create() failed", FX_KV("status", status));
      result.set_err(fuchsia::gpu::agis::Error::INTERNAL_ERROR);
      callback(std::move(result));
      return;
    }

    // Satisfy hanging get for retrieval of the vulkan socket end.
    if (entry->vulkan_socket_callback) {
      fuchsia::gpu::agis::ComponentRegistry_GetVulkanSocket_Result vulkan_socket_result;
      fuchsia::gpu::agis::ComponentRegistry_GetVulkanSocket_Response vulkan_socket_response(
          std::move(entry->vulkan_socket));
      vulkan_socket_result.set_response(std::move(vulkan_socket_response));
      entry->vulkan_socket_callback(std::move(vulkan_socket_result));
      entry->vulkan_socket_callback = nullptr;
    }

    fuchsia::gpu::agis::Connector_GetSocket_Response response(std::move(ffx_socket));
    result.set_response(std::move(response));
    callback(std::move(result));
  }

  void AddBinding(std::unique_ptr<ConnectorImpl> connector,
                  fidl::InterfaceRequest<fuchsia::gpu::agis::Connector> &&request) {
    bindings_.AddBinding(std::move(connector), std::move(request));
  }

 private:
  fidl::BindingSet<fuchsia::gpu::agis::Connector, std::unique_ptr<fuchsia::gpu::agis::Connector>>
      bindings_;
};

int main(int argc, const char **argv) {
  async::Loop loop(&kAsyncLoopConfigAttachToCurrentThread);
  srand(static_cast<unsigned int>(time(nullptr)));
  fuchsia_logging::SetTags({"agis"});

  auto context = sys::ComponentContext::CreateAndServeOutgoingDirectory();
  context->outgoing()->AddPublicService(
      fidl::InterfaceRequestHandler<fuchsia::gpu::agis::ComponentRegistry>(
          [](fidl::InterfaceRequest<fuchsia::gpu::agis::ComponentRegistry> request) {
            auto component_registry = std::make_unique<ComponentRegistryImpl>();
            component_registry->AddBinding(std::move(component_registry), std::move(request));
          }));
  context->outgoing()->AddPublicService(fidl::InterfaceRequestHandler<fuchsia::gpu::agis::Observer>(
      [](fidl::InterfaceRequest<fuchsia::gpu::agis::Observer> request) {
        auto observer = std::make_unique<ObserverImpl>();
        observer->AddBinding(std::move(observer), std::move(request));
      }));
  context->outgoing()->AddPublicService(
      fidl::InterfaceRequestHandler<fuchsia::gpu::agis::Connector>(
          [](fidl::InterfaceRequest<fuchsia::gpu::agis::Connector> request) {
            auto connector = std::make_unique<ConnectorImpl>();
            connector->AddBinding(std::move(connector), std::move(request));
          }));

  return loop.Run();
}