310 lines
13 KiB
C++
310 lines
13 KiB
C++
// Licensed to the Apache Software Foundation (ASF) under one
|
|
// or more contributor license agreements. See the NOTICE file
|
|
// distributed with this work for additional information
|
|
// regarding copyright ownership. The ASF licenses this file
|
|
// to you under the Apache License, Version 2.0 (the
|
|
// "License"); you may not use this file except in compliance
|
|
// with the License. You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing,
|
|
// software distributed under the License is distributed on an
|
|
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
// KIND, either express or implied. See the License for the
|
|
// specific language governing permissions and limitations
|
|
// under the License.
|
|
|
|
#pragma once
|
|
|
|
#include <functional>
|
|
#include <memory>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
#include "arrow/buffer.h"
|
|
#include "arrow/status.h"
|
|
#include "arrow/util/macros.h"
|
|
#include "arrow/util/visibility.h"
|
|
#include "plasma/common.h"
|
|
|
|
using arrow::Buffer;
|
|
using arrow::Status;
|
|
|
|
namespace plasma {
|
|
|
|
/// Object buffer data structure.
|
|
struct ObjectBuffer {
|
|
/// The data buffer.
|
|
std::shared_ptr<Buffer> data;
|
|
/// The metadata buffer.
|
|
std::shared_ptr<Buffer> metadata;
|
|
/// The device number.
|
|
int device_num;
|
|
};
|
|
|
|
class ARROW_EXPORT PlasmaClient {
|
|
public:
|
|
PlasmaClient();
|
|
~PlasmaClient();
|
|
|
|
/// Connect to the local plasma store. Return the resulting connection.
|
|
///
|
|
/// \param store_socket_name The name of the UNIX domain socket to use to
|
|
/// connect to the Plasma store.
|
|
/// \param manager_socket_name The name of the UNIX domain socket to use to
|
|
/// connect to the local Plasma manager. If this is "", then this
|
|
/// function will not connect to a manager.
|
|
/// Note that plasma manager is no longer supported, this function
|
|
/// will return failure if this is not "".
|
|
/// \param release_delay Deprecated (not used).
|
|
/// \param num_retries number of attempts to connect to IPC socket, default 50
|
|
/// \return The return status.
|
|
Status Connect(const std::string& store_socket_name,
|
|
const std::string& manager_socket_name = "", int release_delay = 0,
|
|
int num_retries = -1);
|
|
|
|
/// Set runtime options for this client.
|
|
///
|
|
/// \param client_name The name of the client, used in debug messages.
|
|
/// \param output_memory_quota The memory quota in bytes for objects created by
|
|
/// this client.
|
|
Status SetClientOptions(const std::string& client_name, int64_t output_memory_quota);
|
|
|
|
/// Create an object in the Plasma Store. Any metadata for this object must be
|
|
/// be passed in when the object is created.
|
|
///
|
|
/// \param object_id The ID to use for the newly created object.
|
|
/// \param data_size The size in bytes of the space to be allocated for this
|
|
/// object's
|
|
/// data (this does not include space used for metadata).
|
|
/// \param metadata The object's metadata. If there is no metadata, this
|
|
/// pointer
|
|
/// should be NULL.
|
|
/// \param metadata_size The size in bytes of the metadata. If there is no
|
|
/// metadata, this should be 0.
|
|
/// \param data The address of the newly created object will be written here.
|
|
/// \param device_num The number of the device where the object is being
|
|
/// created.
|
|
/// device_num = 0 corresponds to the host,
|
|
/// device_num = 1 corresponds to GPU0,
|
|
/// device_num = 2 corresponds to GPU1, etc.
|
|
/// \param evict_if_full Whether to evict other objects to make space for
|
|
/// this object.
|
|
/// \return The return status.
|
|
///
|
|
/// The returned object must be released once it is done with. It must also
|
|
/// be either sealed or aborted.
|
|
Status Create(const ObjectID& object_id, int64_t data_size, const uint8_t* metadata,
|
|
int64_t metadata_size, std::shared_ptr<Buffer>* data, int device_num = 0,
|
|
bool evict_if_full = true);
|
|
|
|
/// Create and seal an object in the object store. This is an optimization
|
|
/// which allows small objects to be created quickly with fewer messages to
|
|
/// the store.
|
|
///
|
|
/// \param object_id The ID of the object to create.
|
|
/// \param data The data for the object to create.
|
|
/// \param metadata The metadata for the object to create.
|
|
/// \param evict_if_full Whether to evict other objects to make space for
|
|
/// this object.
|
|
/// \return The return status.
|
|
Status CreateAndSeal(const ObjectID& object_id, const std::string& data,
|
|
const std::string& metadata, bool evict_if_full = true);
|
|
|
|
/// Create and seal multiple objects in the object store. This is an optimization
|
|
/// of CreateAndSeal to eliminate the cost of IPC per object.
|
|
///
|
|
/// \param object_ids The vector of IDs of the objects to create.
|
|
/// \param data The vector of data for the objects to create.
|
|
/// \param metadata The vector of metadata for the objects to create.
|
|
/// \param evict_if_full Whether to evict other objects to make space for
|
|
/// these objects.
|
|
/// \return The return status.
|
|
Status CreateAndSealBatch(const std::vector<ObjectID>& object_ids,
|
|
const std::vector<std::string>& data,
|
|
const std::vector<std::string>& metadata,
|
|
bool evict_if_full = true);
|
|
|
|
/// Get some objects from the Plasma Store. This function will block until the
|
|
/// objects have all been created and sealed in the Plasma Store or the
|
|
/// timeout expires.
|
|
///
|
|
/// If an object was not retrieved, the corresponding metadata and data
|
|
/// fields in the ObjectBuffer structure will evaluate to false.
|
|
/// Objects are automatically released by the client when their buffers
|
|
/// get out of scope.
|
|
///
|
|
/// \param object_ids The IDs of the objects to get.
|
|
/// \param timeout_ms The amount of time in milliseconds to wait before this
|
|
/// request times out. If this value is -1, then no timeout is set.
|
|
/// \param[out] object_buffers The object results.
|
|
/// \return The return status.
|
|
Status Get(const std::vector<ObjectID>& object_ids, int64_t timeout_ms,
|
|
std::vector<ObjectBuffer>* object_buffers);
|
|
|
|
/// Deprecated variant of Get() that doesn't automatically release buffers
|
|
/// when they get out of scope.
|
|
///
|
|
/// \param object_ids The IDs of the objects to get.
|
|
/// \param num_objects The number of object IDs to get.
|
|
/// \param timeout_ms The amount of time in milliseconds to wait before this
|
|
/// request times out. If this value is -1, then no timeout is set.
|
|
/// \param object_buffers An array where the results will be stored.
|
|
/// \return The return status.
|
|
///
|
|
/// The caller is responsible for releasing any retrieved objects, but it
|
|
/// should not release objects that were not retrieved.
|
|
Status Get(const ObjectID* object_ids, int64_t num_objects, int64_t timeout_ms,
|
|
ObjectBuffer* object_buffers);
|
|
|
|
/// Tell Plasma that the client no longer needs the object. This should be
|
|
/// called after Get() or Create() when the client is done with the object.
|
|
/// After this call, the buffer returned by Get() is no longer valid.
|
|
///
|
|
/// \param object_id The ID of the object that is no longer needed.
|
|
/// \return The return status.
|
|
Status Release(const ObjectID& object_id);
|
|
|
|
/// Check if the object store contains a particular object and the object has
|
|
/// been sealed. The result will be stored in has_object.
|
|
///
|
|
/// @todo: We may want to indicate if the object has been created but not
|
|
/// sealed.
|
|
///
|
|
/// \param object_id The ID of the object whose presence we are checking.
|
|
/// \param has_object The function will write true at this address if
|
|
/// the object is present and false if it is not present.
|
|
/// \return The return status.
|
|
Status Contains(const ObjectID& object_id, bool* has_object);
|
|
|
|
/// List all the objects in the object store.
|
|
///
|
|
/// This API is experimental and might change in the future.
|
|
///
|
|
/// \param[out] objects ObjectTable of objects in the store. For each entry
|
|
/// in the map, the following fields are available:
|
|
/// - metadata_size: Size of the object metadata in bytes
|
|
/// - data_size: Size of the object data in bytes
|
|
/// - ref_count: Number of clients referencing the object buffer
|
|
/// - create_time: Unix timestamp of the object creation
|
|
/// - construct_duration: Object creation time in seconds
|
|
/// - state: Is the object still being created or already sealed?
|
|
/// \return The return status.
|
|
Status List(ObjectTable* objects);
|
|
|
|
/// Abort an unsealed object in the object store. If the abort succeeds, then
|
|
/// it will be as if the object was never created at all. The unsealed object
|
|
/// must have only a single reference (the one that would have been removed by
|
|
/// calling Seal).
|
|
///
|
|
/// \param object_id The ID of the object to abort.
|
|
/// \return The return status.
|
|
Status Abort(const ObjectID& object_id);
|
|
|
|
/// Seal an object in the object store. The object will be immutable after
|
|
/// this
|
|
/// call.
|
|
///
|
|
/// \param object_id The ID of the object to seal.
|
|
/// \return The return status.
|
|
Status Seal(const ObjectID& object_id);
|
|
|
|
/// Delete an object from the object store. This currently assumes that the
|
|
/// object is present, has been sealed and not used by another client. Otherwise,
|
|
/// it is a no operation.
|
|
///
|
|
/// \todo We may want to allow the deletion of objects that are not present or
|
|
/// haven't been sealed.
|
|
///
|
|
/// \param object_id The ID of the object to delete.
|
|
/// \return The return status.
|
|
Status Delete(const ObjectID& object_id);
|
|
|
|
/// Delete a list of objects from the object store. This currently assumes that the
|
|
/// object is present, has been sealed and not used by another client. Otherwise,
|
|
/// it is a no operation.
|
|
///
|
|
/// \param object_ids The list of IDs of the objects to delete.
|
|
/// \return The return status. If all the objects are nonexistent, return OK.
|
|
Status Delete(const std::vector<ObjectID>& object_ids);
|
|
|
|
/// Delete objects until we have freed up num_bytes bytes or there are no more
|
|
/// released objects that can be deleted.
|
|
///
|
|
/// \param num_bytes The number of bytes to try to free up.
|
|
/// \param num_bytes_evicted Out parameter for total number of bytes of space
|
|
/// retrieved.
|
|
/// \return The return status.
|
|
Status Evict(int64_t num_bytes, int64_t& num_bytes_evicted);
|
|
|
|
/// Bump objects up in the LRU cache, i.e. treat them as recently accessed.
|
|
/// Objects that do not exist in the store will be ignored.
|
|
///
|
|
/// \param object_ids The IDs of the objects to bump.
|
|
/// \return The return status.
|
|
Status Refresh(const std::vector<ObjectID>& object_ids);
|
|
|
|
/// Compute the hash of an object in the object store.
|
|
///
|
|
/// \param object_id The ID of the object we want to hash.
|
|
/// \param digest A pointer at which to return the hash digest of the object.
|
|
/// The pointer must have at least kDigestSize bytes allocated.
|
|
/// \return The return status.
|
|
Status Hash(const ObjectID& object_id, uint8_t* digest);
|
|
|
|
/// Subscribe to notifications when objects are sealed in the object store.
|
|
/// Whenever an object is sealed, a message will be written to the client
|
|
/// socket that is returned by this method.
|
|
///
|
|
/// \param fd Out parameter for the file descriptor the client should use to
|
|
/// read notifications
|
|
/// from the object store about sealed objects.
|
|
/// \return The return status.
|
|
Status Subscribe(int* fd);
|
|
|
|
/// Receive next object notification for this client if Subscribe has been called.
|
|
///
|
|
/// \param fd The file descriptor we are reading the notification from.
|
|
/// \param object_id Out parameter, the object_id of the object that was sealed.
|
|
/// \param data_size Out parameter, the data size of the object that was sealed.
|
|
/// \param metadata_size Out parameter, the metadata size of the object that was sealed.
|
|
/// \return The return status.
|
|
Status GetNotification(int fd, ObjectID* object_id, int64_t* data_size,
|
|
int64_t* metadata_size);
|
|
|
|
Status DecodeNotifications(const uint8_t* buffer, std::vector<ObjectID>* object_ids,
|
|
std::vector<int64_t>* data_sizes,
|
|
std::vector<int64_t>* metadata_sizes);
|
|
|
|
/// Disconnect from the local plasma instance, including the local store and
|
|
/// manager.
|
|
///
|
|
/// \return The return status.
|
|
Status Disconnect();
|
|
|
|
/// Get the current debug string from the plasma store server.
|
|
///
|
|
/// \return The debug string.
|
|
std::string DebugString();
|
|
|
|
/// Get the memory capacity of the store.
|
|
///
|
|
/// \return Memory capacity of the store in bytes.
|
|
int64_t store_capacity();
|
|
|
|
private:
|
|
friend class PlasmaBuffer;
|
|
friend class PlasmaMutableBuffer;
|
|
FRIEND_TEST(TestPlasmaStore, GetTest);
|
|
FRIEND_TEST(TestPlasmaStore, LegacyGetTest);
|
|
FRIEND_TEST(TestPlasmaStore, AbortTest);
|
|
|
|
bool IsInUse(const ObjectID& object_id);
|
|
|
|
class ARROW_NO_EXPORT Impl;
|
|
std::shared_ptr<Impl> impl_;
|
|
};
|
|
|
|
} // namespace plasma
|