registry.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <string>
 
#include <viam/sdk/common/grpc_fwd.hpp>
#include <viam/sdk/config/resource.hpp>
#include <viam/sdk/resource/resource.hpp>
#include <viam/sdk/resource/resource_api.hpp>
#include <viam/sdk/resource/resource_manager.hpp>
#include <viam/sdk/resource/resource_server_base.hpp>
#include <viam/sdk/rpc/server.hpp>
 
namespace google {
namespace protobuf {
 
class ServiceDescriptor;
 
}
}  // namespace google
 
namespace viam {
namespace sdk {
 
// TODO(RSDK-6617): one class per header
class ResourceServerRegistration {
   public:
    ResourceServerRegistration(const ::google::protobuf::ServiceDescriptor* service_descriptor);
 
    virtual ~ResourceServerRegistration();
 
    virtual std::shared_ptr<ResourceServer> create_resource_server(
        std::shared_ptr<ResourceManager> manager, Server& server) const = 0;
 
    const ::google::protobuf::ServiceDescriptor* service_descriptor() const;
 
   private:
    const ::google::protobuf::ServiceDescriptor* service_descriptor_;
};
 
class ResourceClientRegistration {
   public:
    ResourceClientRegistration() = default;
 
    virtual ~ResourceClientRegistration();
 
    virtual std::shared_ptr<Resource> create_rpc_client(
        std::string name, std::shared_ptr<GrpcChannel> channel) const = 0;
};
 
// TODO(RSDK-6616): instead of std::functions, consider making these functions
// virtual
class ModelRegistration {
   public:
    ModelRegistration(
        API api,
        Model model,
        std::function<std::shared_ptr<Resource>(Dependencies, ResourceConfig)> constructor);
 
    ModelRegistration(
        API api,
        Model model,
        std::function<std::shared_ptr<Resource>(Dependencies, ResourceConfig)> constructor,
        std::function<std::vector<std::string>(ResourceConfig)> validator);
 
    const API& api() const;
    const Model& model() const;
 
    std::function<std::shared_ptr<Resource>(Dependencies, ResourceConfig)> construct_resource;
 
    std::function<std::vector<std::string>(ResourceConfig)> validate;
 
   private:
    // default_validator is the default validator for all models if no validator
    // is provided in construction. No dependencies are returned.
    Model model_;
    API api_;
    static std::vector<std::string> default_validator(ResourceConfig cfg) {
        return {};
    };
};
 
class Registry {
   public:
    static Registry& get();
 
    void register_model(std::shared_ptr<const ModelRegistration> resource);
 
    std::shared_ptr<const ModelRegistration> lookup_model(const std::string& name) const;
 
    std::shared_ptr<const ModelRegistration> lookup_model(const API& api, const Model& model) const;
 
    template <typename ResourceClientT>
    void register_resource_client() {
        class ResourceClientRegistration2 final : public ResourceClientRegistration {
           public:
            using ResourceClientRegistration::ResourceClientRegistration;
 
            std::shared_ptr<Resource> create_rpc_client(
                std::string name, std::shared_ptr<GrpcChannel> chan) const override {
                return std::make_shared<ResourceClientT>(std::move(name), std::move(chan));
            }
        };
 
        Registry::register_resource_client_(API::get<typename ResourceClientT::interface_type>(),
                                            std::make_shared<ResourceClientRegistration2>());
    }
 
    template <typename ResourceServerT>
    void register_resource_server() {
        class ResourceServerRegistration2 final : public ResourceServerRegistration {
           public:
            using ResourceServerRegistration::ResourceServerRegistration;
            std::shared_ptr<ResourceServer> create_resource_server(
                std::shared_ptr<ResourceManager> manager, Server& server) const override {
                auto rs = std::make_shared<ResourceServerT>(manager);
                server.register_service(rs.get());
                return rs;
            }
        };
 
        const google::protobuf::ServiceDescriptor* sd =
            get_service_descriptor_(ResourceServerT::service_type::service_full_name());
        Registry::register_resource_server_(API::get<typename ResourceServerT::interface_type>(),
                                            std::make_shared<ResourceServerRegistration2>(sd));
    }
 
    template <typename ResourceClientT, typename ResourceServerT>
    void register_resource() {
        register_resource_client<ResourceClientT>();
        register_resource_server<ResourceServerT>();
    }
 
    std::shared_ptr<const ResourceServerRegistration> lookup_resource_server(const API& api) const;
 
    std::shared_ptr<const ResourceClientRegistration> lookup_resource_client(const API& api) const;
 
    const std::unordered_map<std::string, std::shared_ptr<const ModelRegistration>>&
    registered_models() const;
 
    const std::unordered_map<API, std::shared_ptr<const ResourceServerRegistration>>&
    registered_resource_servers() const;
 
   private:
    friend class Instance;
    Registry() = default;
 
    void initialize();
 
    mutable std::mutex lock_;
 
    std::unordered_map<std::string, std::shared_ptr<const ModelRegistration>> resources_;
 
    std::unordered_map<API, std::shared_ptr<const ResourceClientRegistration>> client_apis_;
    std::unordered_map<API, std::shared_ptr<const ResourceServerRegistration>> server_apis_;
 
    void register_resources();
 
    void register_resource_server_(
        API api, std::shared_ptr<ResourceServerRegistration> resource_registration);
 
    void register_resource_client_(
        API api, std::shared_ptr<ResourceClientRegistration> resource_registration);
 
    static const google::protobuf::ServiceDescriptor* get_service_descriptor_(
        const char* service_full_name);
 
    std::shared_ptr<const ModelRegistration> lookup_model_inlock_(
        const std::string& name, const std::lock_guard<std::mutex>&) const;
};
 
}  // namespace sdk
}  // namespace viam