Files
ResInsight/GrpcInterface/RiaGrpcKeyValueStoreService.cpp
T

269 lines
12 KiB
C++

/////////////////////////////////////////////////////////////////////////////////
//
// Copyright (C) 2025- Equinor ASA
//
// ResInsight is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// ResInsight is distributed in the hope that it will be useful, but WITHOUT ANY
// WARRANTY; without even the implied warranty of MERCHANTABILITY or
// FITNESS FOR A PARTICULAR PURPOSE.
//
// See the GNU General Public License at <http://www.gnu.org/licenses/gpl.html>
// for more details.
//
//////////////////////////////////////////////////////////////////////////////////
#include "RiaGrpcKeyValueStoreService.h"
#include "KeyValueStore.pb.h"
#include "RiaApplication.h"
#include "RiaGrpcCallbacks.h"
#include "RiaGrpcCaseService.h"
#include "RiaGrpcHelper.h"
#include "RiaKeyValueStoreUtil.h"
#include "RiaLogging.h"
#include "cafAssert.h"
using namespace rips;
#define NUM_CONCURRENT_CLIENT_TO_SERVER_STREAMS 10
class RiaKeyValueStoreStateHandler
{
using Status = grpc::Status;
public:
//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
RiaKeyValueStoreStateHandler( bool clientStreamer = false )
: m_streamedValueCount( 0u )
, m_cellCount( 0u )
{
}
//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
size_t totalValueCount() const { return m_cellCount; }
//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
size_t streamedValueCount() const { return m_streamedValueCount; }
//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
Status init( const std::string& name, size_t numElements )
{
RiaLogging::debug(
QString( "Initializing stream: %1 size: %2" ).arg( QString::fromStdString( name ) ).arg( numElements ).toStdString() );
m_name = name;
m_cellCount = numElements;
return grpc::Status::OK;
}
//--------------------------------------------------------------------------------------------------
/// Client streamers need to be initialised with the encapsulated parameters
//--------------------------------------------------------------------------------------------------
Status init( const KeyValueStoreInputChunk* chunk )
{
CAF_ASSERT( chunk );
if ( chunk->has_parameters() )
{
return init( chunk->parameters().name(), chunk->parameters().num_elements() );
}
return grpc::Status( grpc::INVALID_ARGUMENT, "Need to have parameters in first message" );
}
//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
Status init( const KeyValueStoreOutputRequest* request )
{
std::string name = request->name();
RiaLogging::debug(
QString( "Output stream request. Name='%1'" ).arg( QString::fromStdString( name ) ).toStdString() );
auto result = RiaApplication::instance()->keyValueStore()->get( name );
if ( result.has_value() )
{
m_data = RiaKeyValueStoreUtil::convertToFloatVector( result );
m_name = name;
m_cellCount = m_data.size();
return grpc::Status::OK;
}
else
{
return grpc::Status( grpc::NOT_FOUND, "No matching key found." );
}
}
//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
Status assignStreamReply( KeyValueStoreOutputChunk* reply )
{
// How many data units will fit into one stream package?
const size_t packageSize = RiaGrpcHelper::numberOfDataUnitsInPackage( sizeof( float ) );
const size_t remainingData = m_cellCount - m_streamedValueCount;
const size_t chunkSize = std::min( packageSize, remainingData );
if ( chunkSize == 0 )
{
return grpc::Status( grpc::OUT_OF_RANGE,
"We've reached the end. This is not an error but means transmission is finished" );
}
reply->mutable_values()->Reserve( static_cast<int>( chunkSize ) );
for ( size_t i = 0; i < chunkSize; ++i )
{
reply->add_values( m_data[m_streamedValueCount + i] );
}
m_streamedValueCount += chunkSize;
return grpc::Status::OK;
}
//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
Status receiveStreamRequest( const KeyValueStoreInputChunk* request, ClientToServerStreamReply* reply )
{
CAF_ASSERT( request );
CAF_ASSERT( reply );
if ( request->has_values() )
{
auto values = request->values().values();
if ( !values.empty() )
{
size_t currentCellIdx = m_streamedValueCount;
m_streamedValueCount += values.size();
if ( !values.empty() )
{
m_data.insert( m_data.end(), values.begin(), values.end() );
}
if ( m_streamedValueCount > m_cellCount )
{
return grpc::Status( grpc::OUT_OF_RANGE, "Attempting to write out of bounds" );
}
RiaLogging::debug( QString( "Received stream request. Start index: %1. Size after request: %2" )
.arg( currentCellIdx )
.arg( m_data.size() )
.toStdString() );
reply->set_accepted_value_count( static_cast<int64_t>( m_data.size() ) );
return Status::OK;
}
}
return Status::OK;
}
//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
void finish()
{
RiaLogging::debug( QString( "Stream finished: name=%1 size=%2" )
.arg( QString::fromStdString( m_name ) )
.arg( m_data.size() )
.toStdString() );
if ( !m_name.empty() && !m_data.empty() )
{
RiaApplication::instance()->keyValueStore()->set( m_name, RiaKeyValueStoreUtil::convertToByteVector( m_data ) );
}
}
protected:
size_t m_streamedValueCount;
size_t m_cellCount;
std::string m_name;
std::vector<float> m_data;
};
//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
grpc::Status RiaGrpcKeyValueStoreService::SetValue( grpc::ServerContext* context,
const rips::KeyValueStoreInputChunk* request,
rips::ClientToServerStreamReply* reply,
RiaKeyValueStoreStateHandler* stateHandler )
{
return stateHandler->receiveStreamRequest( request, reply );
}
//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
grpc::Status RiaGrpcKeyValueStoreService::GetValue( grpc::ServerContext* context,
const rips::KeyValueStoreOutputRequest* request,
rips::KeyValueStoreOutputChunk* reply,
RiaKeyValueStoreStateHandler* stateHandler )
{
return stateHandler->assignStreamReply( reply );
}
//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
grpc::Status RiaGrpcKeyValueStoreService::RemoveValue( grpc::ServerContext* context,
const rips::KeyValueStoreRemoveRequest* request,
rips::Empty* reply )
{
std::string name = request->name();
if ( !RiaApplication::instance()->keyValueStore()->exists( name ) )
return grpc::Status( grpc::NOT_FOUND, "No matching key found." );
RiaApplication::instance()->keyValueStore()->remove( name );
return Status::OK;
}
//--------------------------------------------------------------------------------------------------
///
//--------------------------------------------------------------------------------------------------
std::vector<RiaGrpcCallbackInterface*> RiaGrpcKeyValueStoreService::createCallbacks()
{
using Self = RiaGrpcKeyValueStoreService;
std::vector<RiaGrpcCallbackInterface*> callbacks;
callbacks = { new RiaGrpcClientToServerStreamCallback<Self,
KeyValueStoreInputChunk,
ClientToServerStreamReply,
RiaKeyValueStoreStateHandler>( this,
&Self::SetValue,
&Self::RequestSetValue,
new RiaKeyValueStoreStateHandler(
true ) ),
new RiaGrpcUnaryCallback<Self, rips::KeyValueStoreRemoveRequest, rips::Empty>( this,
&Self::RemoveValue,
&Self::RequestRemoveValue ) };
for ( int i = 0; i < NUM_CONCURRENT_CLIENT_TO_SERVER_STREAMS; ++i )
{
callbacks.push_back(
new RiaGrpcServerToClientStreamCallback<Self,
KeyValueStoreOutputRequest,
KeyValueStoreOutputChunk,
RiaKeyValueStoreStateHandler>( this,
&Self::GetValue,
&Self::RequestGetValue,
new RiaKeyValueStoreStateHandler ) );
}
return callbacks;
}
static bool RiaGrpcKeyValueStoreService_init =
RiaGrpcServiceFactory::instance()->registerCreator<RiaGrpcKeyValueStoreService>(
typeid( RiaGrpcKeyValueStoreService ).hash_code() );