// Copyright 2021-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.

//

// Redistribution and use in source and binary forms, with or without

// modification, are permitted provided that the following conditions

// are met:

// * Redistributions of source code must retain the above copyright

// notice, this list of conditions and the following disclaimer.

// * Redistributions in binary form must reproduce the above copyright

// notice, this list of conditions and the following disclaimer in the

// documentation and/or other materials provided with the distribution.

// * Neither the name of NVIDIA CORPORATION nor the names of its

// contributors may be used to endorse or promote products derived

// from this software without specific prior written permission.

//

// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY

// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE

// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR

// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR

// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,

// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,

// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR

// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY

// OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT

// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE

// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#pragma once

#include <boost/interprocess/sync/interprocess_mutex.hpp>

#include <boost/interprocess/sync/interprocess_semaphore.hpp>

#include <boost/interprocess/sync/scoped_lock.hpp>

#include <boost/thread/thread_time.hpp>

#include <cstddef>

#include "shm_manager.h"

namespace triton { namespace backend { namespace python {

namespace bi = boost::interprocess;

/// Struct holding the represenation of a message queue inside the shared

/// memory.

/// \param size Total size of the message queue.

/// \param mutex Handle of the mutex variable protecting index.

/// \param index Used element index.

/// \param sem_empty Semaphore object counting the number of empty buffer slots.

/// \param sem_full Semaphore object counting the number of used buffer slots.

struct MessageQueueShm {

bi::interprocess_semaphore sem_empty{0};

bi::interprocess_semaphore sem_full{0};

bi::interprocess_mutex mutex;

std::size_t size;

bi::managed_external_buffer::handle_t buffer;

int head;

int tail;

};

template <typename T>

class MessageQueue {

public:

/// Create a new MessageQueue in the shared memory.

static std::unique_ptr<MessageQueue<T>> Create(

std::unique_ptr<SharedMemoryManager>& shm_pool,

uint32_t message_queue_size)

{

AllocatedSharedMemory<MessageQueueShm> mq_shm =

shm_pool->Construct<MessageQueueShm>();

mq_shm.data_->size = message_queue_size;

AllocatedSharedMemory<T> mq_buffer_shm =

shm_pool->Construct<T>(message_queue_size /* count */);

mq_shm.data_->buffer = mq_buffer_shm.handle_;

mq_shm.data_->head = 0;

mq_shm.data_->tail = 0;

new (&(mq_shm.data_->mutex)) bi::interprocess_mutex{};

new (&(mq_shm.data_->sem_empty))

bi::interprocess_semaphore{message_queue_size};

new (&(mq_shm.data_->sem_full)) bi::interprocess_semaphore{0};

return std::unique_ptr<MessageQueue<T>>(

new MessageQueue<T>(mq_shm, mq_buffer_shm));

}

/// Load an already existing message queue from the shared memory.

static std::unique_ptr<MessageQueue<T>> LoadFromSharedMemory(

std::unique_ptr<SharedMemoryManager>& shm_pool,

bi::managed_external_buffer::handle_t message_queue_handle)

{

AllocatedSharedMemory<MessageQueueShm> mq_shm =

shm_pool->Load<MessageQueueShm>(message_queue_handle);

AllocatedSharedMemory<T> mq_shm_buffer =

shm_pool->Load<T>(mq_shm.data_->buffer);

return std::unique_ptr<MessageQueue<T>>(

new MessageQueue(mq_shm, mq_shm_buffer));

}

/// Push a message inside the message queue.

/// \param message The shared memory handle of the message.

void Push(T message)

{

while (true) {

try {

SemEmptyMutable()->wait();

break;

}

catch (bi::interprocess_exception& ex) {

}

}

{

bi::scoped_lock<bi::interprocess_mutex> lock{*MutexMutable()};

Buffer()[Head()] = message;

HeadIncrement();

}

SemFullMutable()->post();

}

void Push(T message, int const& duration, bool& success)

{

boost::system_time timeout =

boost::get_system_time() + boost::posix_time::milliseconds(duration);

while (true) {

try {

if (!SemEmptyMutable()->timed_wait(timeout)) {

success = false;

return;

} else {

break;

}

}

catch (bi::interprocess_exception& ex) {

}

}

{

timeout =

boost::get_system_time() + boost::posix_time::milliseconds(duration);

bi::scoped_lock<bi::interprocess_mutex> lock{*MutexMutable(), timeout};

if (!lock) {

SemEmptyMutable()->post();

success = false;

return;

}

success = true;

Buffer()[Head()] = message;

HeadIncrement();

}

SemFullMutable()->post();

}

/// Pop a message from the message queue. This call will block until there

/// is a message inside the message queue to return.

/// \return the handle of the new message.

T Pop()

{

T message;

while (true) {

try {

SemFullMutable()->wait();

break;

}

catch (bi::interprocess_exception& ex) {

}

}

{

bi::scoped_lock<bi::interprocess_mutex> lock{*MutexMutable()};

message = Buffer()[Tail()];

TailIncrement();

}

SemEmptyMutable()->post();

return message;

}

T Pop(int const& duration, bool& success)

{

T message = 0;

boost::system_time timeout =

boost::get_system_time() + boost::posix_time::milliseconds(duration);

while (true) {

try {

if (!SemFullMutable()->timed_wait(timeout)) {

success = false;

return message;

} else {

break;

}

}

catch (bi::interprocess_exception& ex) {

}

}

{

timeout =

boost::get_system_time() + boost::posix_time::milliseconds(duration);

bi::scoped_lock<bi::interprocess_mutex> lock{*MutexMutable(), timeout};

if (!lock) {

SemFullMutable()->post();

success = false;

return message;

}

success = true;

message = Buffer()[Tail()];

TailIncrement();

}

SemEmptyMutable()->post();

return message;

}

/// Resets the semaphores for the message queue. This function is useful for

/// when the stub process may have exited unexpectedly and the semaphores need

/// to be restarted so that the message queue is in a proper state.

void ResetSemaphores()

{

new (SemFullMutable()) bi::interprocess_semaphore(0);

new (SemEmptyMutable()) bi::interprocess_semaphore(Size());

new (MutexMutable()) bi::interprocess_mutex;

mq_shm_ptr_->tail = 0;

mq_shm_ptr_->head = 0;

}

/// Get the shared memory handle of MessageQueue

bi::managed_external_buffer::handle_t ShmHandle() { return mq_handle_; }

/// Release the ownership of this object in shared memory.

void Release()

{

if (mq_shm_.data_ != nullptr) {

mq_shm_.data_.release();

}

if (mq_buffer_shm_.data_ != nullptr) {

mq_buffer_shm_.data_.release();

}

}

private:

std::size_t& Size() { return mq_shm_ptr_->size; }

const bi::interprocess_mutex& Mutex() { return mq_shm_ptr_->mutex; }

bi::interprocess_mutex* MutexMutable() { return &(mq_shm_ptr_->mutex); }

int& Head() { return mq_shm_ptr_->head; }

int& Tail() { return mq_shm_ptr_->tail; }

T* Buffer() { return mq_buffer_shm_ptr_; }

const bi::interprocess_semaphore& SemEmpty()

{

return mq_shm_ptr_->sem_empty;

}

bi::interprocess_semaphore* SemEmptyMutable()

{

return &(mq_shm_ptr_->sem_empty);

}

const bi::interprocess_semaphore& SemFull() { return mq_shm_ptr_->sem_full; }

bi::interprocess_semaphore* SemFullMutable()

{

return &(mq_shm_ptr_->sem_full);

}

void HeadIncrement() { mq_shm_ptr_->head = (mq_shm_ptr_->head + 1) % Size(); }

void TailIncrement() { mq_shm_ptr_->tail = (mq_shm_ptr_->tail + 1) % Size(); }

AllocatedSharedMemory<MessageQueueShm> mq_shm_;

AllocatedSharedMemory<T> mq_buffer_shm_;

MessageQueueShm* mq_shm_ptr_;

T* mq_buffer_shm_ptr_;

bi::managed_external_buffer::handle_t mq_handle_;

/// Create/load a Message queue.

/// \param mq_shm Message queue representation in shared memory.

MessageQueue(

AllocatedSharedMemory<MessageQueueShm>& mq_shm,

AllocatedSharedMemory<T>& mq_buffer_shm)

: mq_shm_(std::move(mq_shm)), mq_buffer_shm_(std::move(mq_buffer_shm))

{

mq_buffer_shm_ptr_ = mq_buffer_shm_.data_.get();

mq_shm_ptr_ = mq_shm_.data_.get();

mq_handle_ = mq_shm_.handle_;

}

};

}}} // namespace triton::backend::python