2020-02-13 03:05:15 +01:00
# include "lokimq.h"
2020-03-13 18:02:54 +01:00
# include "lokimq-internal.h"
2020-11-10 22:47:31 +01:00
# include "zmq.hpp"
2020-02-03 03:39:26 +01:00
# include <map>
2020-02-11 07:30:07 +01:00
# include <random>
2020-06-23 00:10:04 +02:00
# include <ostream>
2020-02-11 07:30:07 +01:00
2020-02-03 03:39:26 +01:00
extern " C " {
2020-06-22 21:23:04 +02:00
# include <sodium/core.h>
# include <sodium/crypto_box.h>
# include <sodium/crypto_scalarmult.h>
2020-02-03 03:39:26 +01:00
}
# include "hex.h"
namespace lokimq {
namespace {
/// Creates a message by bt-serializing the given value (string, number, list, or dict)
template < typename T >
zmq : : message_t create_bt_message ( T & & data ) { return create_message ( bt_serialize ( std : : forward < T > ( data ) ) ) ; }
template < typename MessageContainer >
std : : vector < std : : string > as_strings ( const MessageContainer & msgs ) {
std : : vector < std : : string > result ;
result . reserve ( msgs . size ( ) ) ;
for ( const auto & msg : msgs )
result . emplace_back ( msg . template data < char > ( ) , msg . size ( ) ) ;
return result ;
}
2020-02-11 07:30:07 +01:00
void check_not_started ( const std : : thread & proxy_thread , const std : : string & verb ) {
2020-02-06 01:21:27 +01:00
if ( proxy_thread . joinable ( ) )
2020-02-11 07:30:07 +01:00
throw std : : logic_error ( " Cannot " + verb + " after calling `start()` " ) ;
2020-02-03 03:39:26 +01:00
}
2020-03-13 18:02:54 +01:00
} // anonymous namespace
2020-02-03 03:39:26 +01:00
2020-03-13 18:02:54 +01:00
namespace detail {
// Sends a control messages between proxy and threads or between proxy and workers consisting of a
// single command codes with an optional data part (the data frame is omitted if empty).
2020-05-12 20:33:59 +02:00
void send_control ( zmq : : socket_t & sock , std : : string_view cmd , std : : string data ) {
2020-03-13 18:02:54 +01:00
auto c = create_message ( std : : move ( cmd ) ) ;
if ( data . empty ( ) ) {
sock . send ( c , zmq : : send_flags : : none ) ;
} else {
auto d = create_message ( std : : move ( data ) ) ;
sock . send ( c , zmq : : send_flags : : sndmore ) ;
sock . send ( d , zmq : : send_flags : : none ) ;
2020-02-03 03:39:26 +01:00
}
}
1.1.0: invocation-time SN auth; failure responses
This replaces the recognition of SN status to be checked per-command
invocation rather than on connection. As this breaks the API quite
substantially, though doesn't really affect the functionality, it seems
suitable to bump the minor version.
This requires a fundamental shift in how the calling application tells
LokiMQ about service nodes: rather than using a callback invoked on
connection, the application now has to call set_active_sns() (or the
more efficient update_active_sns(), if changes are readily available) to
update the list whenever it changes. LokiMQ then keeps this list
internally and uses it when determining whether to invoke.
This release also brings better request responses on errors: when a
request fails, the data argument will now be set to the failure reason,
one of:
- TIMEOUT
- UNKNOWNCOMMAND
- NOT_A_SERVICE_NODE (the remote isn't running in SN mode)
- FORBIDDEN (auth level denies the request)
- FORBIDDEN_SN (SN required and the remote doesn't see us as a SN)
Some of these (UNKNOWNCOMMAND, NOT_A_SERVICE_NODE, FORBIDDEN) were
already sent by remotes, but there was no connection to a request and so
they would log a warning, but the request would have to time out.
These errors (minus TIMEOUT, plus NO_REPLY_TAG signalling that a command
is a request but didn't include a reply tag) are also sent in response
to regular commands, but they simply result in a log warning showing the
error type and the command that caused the failure when received.
2020-04-13 00:57:19 +02:00
/// Extracts a pubkey and and auth level from a zmq message received on a *listening* socket.
std : : pair < std : : string , AuthLevel > extract_metadata ( zmq : : message_t & msg ) {
auto result = std : : make_pair ( " " s , AuthLevel : : none ) ;
2020-02-28 05:16:43 +01:00
try {
2020-05-12 20:33:59 +02:00
std : : string_view pubkey_hex { msg . gets ( " User-Id " ) } ;
2020-02-28 05:16:43 +01:00
if ( pubkey_hex . size ( ) ! = 64 )
throw std : : logic_error ( " bad user-id " ) ;
assert ( is_hex ( pubkey_hex . begin ( ) , pubkey_hex . end ( ) ) ) ;
1.1.0: invocation-time SN auth; failure responses
This replaces the recognition of SN status to be checked per-command
invocation rather than on connection. As this breaks the API quite
substantially, though doesn't really affect the functionality, it seems
suitable to bump the minor version.
This requires a fundamental shift in how the calling application tells
LokiMQ about service nodes: rather than using a callback invoked on
connection, the application now has to call set_active_sns() (or the
more efficient update_active_sns(), if changes are readily available) to
update the list whenever it changes. LokiMQ then keeps this list
internally and uses it when determining whether to invoke.
This release also brings better request responses on errors: when a
request fails, the data argument will now be set to the failure reason,
one of:
- TIMEOUT
- UNKNOWNCOMMAND
- NOT_A_SERVICE_NODE (the remote isn't running in SN mode)
- FORBIDDEN (auth level denies the request)
- FORBIDDEN_SN (SN required and the remote doesn't see us as a SN)
Some of these (UNKNOWNCOMMAND, NOT_A_SERVICE_NODE, FORBIDDEN) were
already sent by remotes, but there was no connection to a request and so
they would log a warning, but the request would have to time out.
These errors (minus TIMEOUT, plus NO_REPLY_TAG signalling that a command
is a request but didn't include a reply tag) are also sent in response
to regular commands, but they simply result in a log warning showing the
error type and the command that caused the failure when received.
2020-04-13 00:57:19 +02:00
result . first . resize ( 32 , 0 ) ;
from_hex ( pubkey_hex . begin ( ) , pubkey_hex . end ( ) , result . first . begin ( ) ) ;
2020-02-28 05:16:43 +01:00
} catch ( . . . ) { }
try {
1.1.0: invocation-time SN auth; failure responses
This replaces the recognition of SN status to be checked per-command
invocation rather than on connection. As this breaks the API quite
substantially, though doesn't really affect the functionality, it seems
suitable to bump the minor version.
This requires a fundamental shift in how the calling application tells
LokiMQ about service nodes: rather than using a callback invoked on
connection, the application now has to call set_active_sns() (or the
more efficient update_active_sns(), if changes are readily available) to
update the list whenever it changes. LokiMQ then keeps this list
internally and uses it when determining whether to invoke.
This release also brings better request responses on errors: when a
request fails, the data argument will now be set to the failure reason,
one of:
- TIMEOUT
- UNKNOWNCOMMAND
- NOT_A_SERVICE_NODE (the remote isn't running in SN mode)
- FORBIDDEN (auth level denies the request)
- FORBIDDEN_SN (SN required and the remote doesn't see us as a SN)
Some of these (UNKNOWNCOMMAND, NOT_A_SERVICE_NODE, FORBIDDEN) were
already sent by remotes, but there was no connection to a request and so
they would log a warning, but the request would have to time out.
These errors (minus TIMEOUT, plus NO_REPLY_TAG signalling that a command
is a request but didn't include a reply tag) are also sent in response
to regular commands, but they simply result in a log warning showing the
error type and the command that caused the failure when received.
2020-04-13 00:57:19 +02:00
result . second = auth_from_string ( msg . gets ( " X-AuthLevel " ) ) ;
2020-02-28 05:16:43 +01:00
} catch ( . . . ) { }
return result ;
2020-02-03 03:39:26 +01:00
}
} // namespace detail
2020-11-10 22:47:31 +01:00
void LokiMQ : : set_zmq_context_option ( zmq : : ctxopt option , int value ) {
context . set ( option , value ) ;
2020-04-03 01:52:48 +02:00
}
2020-02-03 03:39:26 +01:00
void LokiMQ : : log_level ( LogLevel level ) {
log_lvl . store ( level , std : : memory_order_relaxed ) ;
}
LogLevel LokiMQ : : log_level ( ) const {
return log_lvl . load ( std : : memory_order_relaxed ) ;
}
Add a category command helper class
This allows simplifying:
lmq.add_category("foo", ...);
lmq.add_command("foo", "a", ...);
lmq.add_command("foo", "b", ...);
lmq.add_request_command("foo", "c", ...);
to:
lmq.add_category("foo", ...)
.add_command("a", ...)
.add_command("b", ...)
.add_request_command("b", ...)
;
2020-03-02 20:11:54 +01:00
CatHelper LokiMQ : : add_category ( std : : string name , Access access_level , unsigned int reserved_threads , int max_queue ) {
2020-02-11 07:30:07 +01:00
check_not_started ( proxy_thread , " add a category " ) ;
2020-02-03 03:39:26 +01:00
if ( name . size ( ) > MAX_CATEGORY_LENGTH )
throw std : : runtime_error ( " Invalid category name ` " + name + " ': name too long (> " + std : : to_string ( MAX_CATEGORY_LENGTH ) + " ) " ) ;
if ( name . empty ( ) | | name . find ( ' . ' ) ! = std : : string : : npos )
throw std : : runtime_error ( " Invalid category name ` " + name + " ' " ) ;
auto it = categories . find ( name ) ;
if ( it ! = categories . end ( ) )
throw std : : runtime_error ( " Unable to add category ` " + name + " ': that category already exists " ) ;
Add a category command helper class
This allows simplifying:
lmq.add_category("foo", ...);
lmq.add_command("foo", "a", ...);
lmq.add_command("foo", "b", ...);
lmq.add_request_command("foo", "c", ...);
to:
lmq.add_category("foo", ...)
.add_command("a", ...)
.add_command("b", ...)
.add_request_command("b", ...)
;
2020-03-02 20:11:54 +01:00
CatHelper ret { * this , name } ;
2020-02-03 03:39:26 +01:00
categories . emplace ( std : : move ( name ) , category { access_level , reserved_threads , max_queue } ) ;
Add a category command helper class
This allows simplifying:
lmq.add_category("foo", ...);
lmq.add_command("foo", "a", ...);
lmq.add_command("foo", "b", ...);
lmq.add_request_command("foo", "c", ...);
to:
lmq.add_category("foo", ...)
.add_command("a", ...)
.add_command("b", ...)
.add_request_command("b", ...)
;
2020-03-02 20:11:54 +01:00
return ret ;
2020-02-03 03:39:26 +01:00
}
void LokiMQ : : add_command ( const std : : string & category , std : : string name , CommandCallback callback ) {
2020-02-11 07:30:07 +01:00
check_not_started ( proxy_thread , " add a command " ) ;
2020-02-03 03:39:26 +01:00
if ( name . size ( ) > MAX_COMMAND_LENGTH )
throw std : : runtime_error ( " Invalid command name ` " + name + " ': name too long (> " + std : : to_string ( MAX_COMMAND_LENGTH ) + " ) " ) ;
auto catit = categories . find ( category ) ;
if ( catit = = categories . end ( ) )
throw std : : runtime_error ( " Cannot add a command to unknown category ` " + category + " ' " ) ;
std : : string fullname = category + ' . ' + name ;
if ( command_aliases . count ( fullname ) )
throw std : : runtime_error ( " Cannot add command ` " + fullname + " ': a command alias with that name is already defined " ) ;
2020-02-11 07:30:07 +01:00
auto ins = catit - > second . commands . insert ( { std : : move ( name ) , { std : : move ( callback ) , false } } ) ;
2020-02-03 03:39:26 +01:00
if ( ! ins . second )
throw std : : runtime_error ( " Cannot add command ` " + fullname + " ': that command already exists " ) ;
}
2020-02-11 07:30:07 +01:00
void LokiMQ : : add_request_command ( const std : : string & category , std : : string name , CommandCallback callback ) {
add_command ( category , name , std : : move ( callback ) ) ;
categories . at ( category ) . commands . at ( name ) . second = true ;
}
2020-02-03 03:39:26 +01:00
void LokiMQ : : add_command_alias ( std : : string from , std : : string to ) {
2020-02-11 07:30:07 +01:00
check_not_started ( proxy_thread , " add a command alias " ) ;
2020-02-03 03:39:26 +01:00
if ( from . empty ( ) )
throw std : : runtime_error ( " Cannot add an alias for empty command " ) ;
size_t fromdot = from . find ( ' . ' ) ;
if ( fromdot = = 0 ) // We don't have to have a ., but if we do it can't be at the beginning.
throw std : : runtime_error ( " Invalid command alias ` " + from + " ' " ) ;
size_t todot = to . find ( ' . ' ) ;
if ( todot = = 0 | | todot = = std : : string : : npos ) // must have a dot for the target
throw std : : runtime_error ( " Invalid command alias target ` " + to + " ' " ) ;
if ( fromdot ! = std : : string : : npos ) {
auto catit = categories . find ( from . substr ( 0 , fromdot ) ) ;
if ( catit ! = categories . end ( ) & & catit - > second . commands . count ( from . substr ( fromdot + 1 ) ) )
throw std : : runtime_error ( " Invalid command alias: ` " + from + " ' would mask an existing command " ) ;
}
auto ins = command_aliases . emplace ( std : : move ( from ) , std : : move ( to ) ) ;
if ( ! ins . second )
throw std : : runtime_error ( " Cannot add command alias ` " + ins . first - > first + " ': that alias already exists " ) ;
}
std : : atomic < int > next_id { 1 } ;
/// Accesses a thread-local command socket connected to the proxy's command socket used to issue
/// commands in a thread-safe manner. A mutex is only required here the first time a thread
/// accesses the control socket.
zmq : : socket_t & LokiMQ : : get_control_socket ( ) {
2020-02-11 07:30:07 +01:00
assert ( proxy_thread . joinable ( ) ) ;
2020-02-03 03:39:26 +01:00
// Maps the LokiMQ unique ID to a local thread command socket.
static thread_local std : : map < int , std : : shared_ptr < zmq : : socket_t > > control_sockets ;
static thread_local std : : pair < int , std : : shared_ptr < zmq : : socket_t > > last { - 1 , nullptr } ;
// Optimize by caching the last value; LokiMQ is often a singleton and in that case we're
// going to *always* hit this optimization. Even if it isn't, we're probably likely to need the
// same control socket from the same thread multiple times sequentially so this may still help.
if ( object_id = = last . first )
return * last . second ;
auto it = control_sockets . find ( object_id ) ;
if ( it ! = control_sockets . end ( ) ) {
last = * it ;
return * last . second ;
}
2020-05-12 20:59:22 +02:00
std : : lock_guard lock { control_sockets_mutex } ;
2020-02-03 03:39:26 +01:00
if ( proxy_shutting_down )
throw std : : runtime_error ( " Unable to obtain LokiMQ control socket: proxy thread is shutting down " ) ;
auto control = std : : make_shared < zmq : : socket_t > ( context , zmq : : socket_type : : dealer ) ;
2020-11-10 22:47:31 +01:00
control - > set ( zmq : : sockopt : : linger , 0 ) ;
2020-02-03 03:39:26 +01:00
control - > connect ( SN_ADDR_COMMAND ) ;
thread_control_sockets . push_back ( control ) ;
control_sockets . emplace ( object_id , control ) ;
last . first = object_id ;
last . second = std : : move ( control ) ;
return * last . second ;
}
LokiMQ : : LokiMQ (
std : : string pubkey_ ,
std : : string privkey_ ,
bool service_node ,
SNRemoteAddress lookup ,
2020-03-29 20:20:03 +02:00
Logger logger ,
LogLevel level )
2020-02-03 03:39:26 +01:00
: object_id { next_id + + } , pubkey { std : : move ( pubkey_ ) } , privkey { std : : move ( privkey_ ) } , local_service_node { service_node } ,
2020-03-29 20:20:03 +02:00
sn_lookup { std : : move ( lookup ) } , log_lvl { level } , logger { std : : move ( logger ) }
2020-02-28 05:16:43 +01:00
{
2020-02-03 03:39:26 +01:00
2020-04-26 17:12:04 +02:00
LMQ_TRACE ( " Constructing LokiMQ, id= " , object_id , " , this= " , this ) ;
2020-02-03 03:39:26 +01:00
2020-05-15 06:33:02 +02:00
if ( sodium_init ( ) = = - 1 )
throw std : : runtime_error { " libsodium initialization failed " } ;
2020-02-03 03:39:26 +01:00
if ( pubkey . empty ( ) ! = privkey . empty ( ) ) {
throw std : : invalid_argument ( " LokiMQ construction failed: one (and only one) of pubkey/privkey is empty. Both must be specified, or both empty to generate a key. " ) ;
} else if ( pubkey . empty ( ) ) {
if ( service_node )
throw std : : invalid_argument ( " Cannot construct a service node mode LokiMQ without a keypair " ) ;
LMQ_LOG ( debug , " generating x25519 keypair for remote-only LokiMQ instance " ) ;
pubkey . resize ( crypto_box_PUBLICKEYBYTES ) ;
privkey . resize ( crypto_box_SECRETKEYBYTES ) ;
crypto_box_keypair ( reinterpret_cast < unsigned char * > ( & pubkey [ 0 ] ) , reinterpret_cast < unsigned char * > ( & privkey [ 0 ] ) ) ;
} else if ( pubkey . size ( ) ! = crypto_box_PUBLICKEYBYTES ) {
throw std : : invalid_argument ( " pubkey has invalid size " + std : : to_string ( pubkey . size ( ) ) + " , expected " + std : : to_string ( crypto_box_PUBLICKEYBYTES ) ) ;
} else if ( privkey . size ( ) ! = crypto_box_SECRETKEYBYTES ) {
throw std : : invalid_argument ( " privkey has invalid size " + std : : to_string ( privkey . size ( ) ) + " , expected " + std : : to_string ( crypto_box_SECRETKEYBYTES ) ) ;
} else {
// Verify the pubkey. We could get by with taking just the privkey and just generate this
// for ourselves, but this provides an extra check to make sure we and the caller agree
// cryptographically (e.g. to make sure they don't pass us an ed25519 keypair by mistake)
2020-02-13 05:53:43 +01:00
std : : string verify_pubkey ( crypto_box_PUBLICKEYBYTES , 0 ) ;
2020-02-13 05:39:00 +01:00
crypto_scalarmult_base ( reinterpret_cast < unsigned char * > ( & verify_pubkey [ 0 ] ) , reinterpret_cast < unsigned char * > ( & privkey [ 0 ] ) ) ;
2020-02-03 03:39:26 +01:00
if ( verify_pubkey ! = pubkey )
throw std : : invalid_argument ( " Invalid pubkey/privkey values given to LokiMQ construction: pubkey verification failed " ) ;
}
}
void LokiMQ : : start ( ) {
2020-02-06 01:21:27 +01:00
if ( proxy_thread . joinable ( ) )
2020-02-03 03:39:26 +01:00
throw std : : logic_error ( " Cannot call start() multiple times! " ) ;
2020-02-28 05:16:43 +01:00
// If we're not binding to anything then we don't listen, i.e. we can only establish outbound
// connections. Don't allow this if we are in service_node mode because, if we aren't
// listening, we are useless as a service node.
if ( bind . empty ( ) & & local_service_node )
throw std : : invalid_argument { " Cannot create a service node listener with no address(es) to bind " } ;
2020-02-03 03:39:26 +01:00
LMQ_LOG ( info , " Initializing LokiMQ " , bind . empty ( ) ? " remote-only " : " listener " , " with pubkey " , to_hex ( pubkey ) ) ;
2020-11-10 22:47:31 +01:00
int zmq_socket_limit = context . get ( zmq : : ctxopt : : socket_limit ) ;
2020-04-17 21:13:04 +02:00
if ( MAX_SOCKETS > 1 & & MAX_SOCKETS < = zmq_socket_limit )
2020-11-10 22:47:31 +01:00
context . set ( zmq : : ctxopt : : max_sockets , MAX_SOCKETS ) ;
2020-04-17 21:13:04 +02:00
else
LMQ_LOG ( error , " Not applying LokiMQ::MAX_SOCKETS setting: " , MAX_SOCKETS , " must be in [1, " , zmq_socket_limit , " ] " ) ;
2020-02-03 03:39:26 +01:00
// We bind `command` here so that the `get_control_socket()` below is always connecting to a
// bound socket, but we do nothing else here: the proxy thread is responsible for everything
// except binding it.
command . bind ( SN_ADDR_COMMAND ) ;
proxy_thread = std : : thread { & LokiMQ : : proxy_loop , this } ;
LMQ_LOG ( debug , " Waiting for proxy thread to get ready... " ) ;
auto & control = get_control_socket ( ) ;
detail : : send_control ( control , " START " ) ;
2020-02-13 03:10:40 +01:00
LMQ_TRACE ( " Sent START command " ) ;
2020-02-03 03:39:26 +01:00
zmq : : message_t ready_msg ;
std : : vector < zmq : : message_t > parts ;
2020-02-06 05:50:31 +01:00
try { recv_message_parts ( control , parts ) ; }
2020-02-03 03:39:26 +01:00
catch ( const zmq : : error_t & e ) { throw std : : runtime_error ( " Failure reading from LokiMQ::Proxy thread: " s + e . what ( ) ) ; }
if ( ! ( parts . size ( ) = = 1 & & view ( parts . front ( ) ) = = " READY " ) )
throw std : : runtime_error ( " Invalid startup message from proxy thread (didn't get expected READY message) " ) ;
LMQ_LOG ( debug , " Proxy thread is ready " ) ;
}
2020-02-28 05:16:43 +01:00
void LokiMQ : : listen_curve ( std : : string bind_addr , AllowFunc allow_connection ) {
// TODO: there's no particular reason we can't start listening after starting up; just needs to
// be implemented. (But if we can start we'll probably also want to be able to stop, so it's
// more than just binding that needs implementing).
check_not_started ( proxy_thread , " start listening " ) ;
bind . emplace_back ( std : : move ( bind_addr ) , bind_data { true , std : : move ( allow_connection ) } ) ;
}
void LokiMQ : : listen_plain ( std : : string bind_addr , AllowFunc allow_connection ) {
// TODO: As above.
check_not_started ( proxy_thread , " start listening " ) ;
bind . emplace_back ( std : : move ( bind_addr ) , bind_data { false , std : : move ( allow_connection ) } ) ;
}
2020-02-03 03:39:26 +01:00
2020-02-11 07:30:07 +01:00
std : : pair < LokiMQ : : category * , const std : : pair < LokiMQ : : CommandCallback , bool > * > LokiMQ : : get_command ( std : : string & command ) {
2020-02-03 03:39:26 +01:00
if ( command . size ( ) > MAX_CATEGORY_LENGTH + 1 + MAX_COMMAND_LENGTH ) {
LMQ_LOG ( warn , " Invalid command ' " , command , " ': command too long " ) ;
return { } ;
}
if ( ! command_aliases . empty ( ) ) {
auto it = command_aliases . find ( command ) ;
if ( it ! = command_aliases . end ( ) )
command = it - > second ;
}
auto dot = command . find ( ' . ' ) ;
if ( dot = = 0 | | dot = = std : : string : : npos ) {
LMQ_LOG ( warn , " Invalid command ' " , command , " ': expected <category>.<command> " ) ;
return { } ;
}
2020-02-25 03:20:56 +01:00
std : : string catname = command . substr ( 0 , dot ) ;
2020-02-03 03:39:26 +01:00
std : : string cmd = command . substr ( dot + 1 ) ;
auto catit = categories . find ( catname ) ;
if ( catit = = categories . end ( ) ) {
LMQ_LOG ( warn , " Invalid command category ' " , catname , " ' " ) ;
return { } ;
}
const auto & category = catit - > second ;
auto callback_it = category . commands . find ( cmd ) ;
if ( callback_it = = category . commands . end ( ) ) {
LMQ_LOG ( warn , " Invalid command ' " , command , " ' " ) ;
return { } ;
}
return { & catit - > second , & callback_it - > second } ;
}
2020-02-28 22:54:00 +01:00
void LokiMQ : : set_batch_threads ( int threads ) {
2020-02-06 01:21:27 +01:00
if ( proxy_thread . joinable ( ) )
throw std : : logic_error ( " Cannot change reserved batch threads after calling `start()` " ) ;
2020-02-28 22:54:00 +01:00
if ( threads < - 1 ) // -1 is the default which is based on general threads
throw std : : out_of_range ( " Invalid set_batch_threads() value " + std : : to_string ( threads ) ) ;
2020-02-06 01:21:27 +01:00
batch_jobs_reserved = threads ;
}
2020-02-28 22:54:00 +01:00
void LokiMQ : : set_reply_threads ( int threads ) {
if ( proxy_thread . joinable ( ) )
throw std : : logic_error ( " Cannot change reserved reply threads after calling `start()` " ) ;
if ( threads < - 1 ) // -1 is the default which is based on general threads
throw std : : out_of_range ( " Invalid set_reply_threads() value " + std : : to_string ( threads ) ) ;
reply_jobs_reserved = threads ;
}
void LokiMQ : : set_general_threads ( int threads ) {
2020-02-06 01:21:27 +01:00
if ( proxy_thread . joinable ( ) )
throw std : : logic_error ( " Cannot change general thread count after calling `start()` " ) ;
2020-02-28 22:54:00 +01:00
if ( threads < 1 )
throw std : : out_of_range ( " Invalid set_general_threads() value " + std : : to_string ( threads ) + " : general threads must be > 0 " ) ;
2020-02-06 01:21:27 +01:00
general_workers = threads ;
}
2020-04-24 23:58:24 +02:00
LokiMQ : : run_info & LokiMQ : : run_info : : load ( category * cat_ , std : : string command_ , ConnectionID conn_ , Access access_ , std : : string remote_ ,
2020-02-28 22:54:00 +01:00
std : : vector < zmq : : message_t > data_parts_ , const std : : pair < CommandCallback , bool > * callback_ ) {
2020-06-08 02:30:46 +02:00
reset ( ) ;
2020-02-28 22:54:00 +01:00
cat = cat_ ;
command = std : : move ( command_ ) ;
conn = std : : move ( conn_ ) ;
2020-04-24 02:52:39 +02:00
access = std : : move ( access_ ) ;
2020-04-24 23:58:24 +02:00
remote = std : : move ( remote_ ) ;
2020-02-28 22:54:00 +01:00
data_parts = std : : move ( data_parts_ ) ;
2020-06-08 02:30:46 +02:00
to_run = callback_ ;
return * this ;
}
LokiMQ : : run_info & LokiMQ : : run_info : : load ( category * cat_ , std : : string command_ , std : : string remote_ , std : : function < void ( ) > callback ) {
reset ( ) ;
is_injected = true ;
cat = cat_ ;
command = std : : move ( command_ ) ;
conn = { } ;
access = { } ;
remote = std : : move ( remote_ ) ;
to_run = std : : move ( callback ) ;
2020-02-06 01:21:27 +01:00
return * this ;
}
2020-02-28 22:54:00 +01:00
LokiMQ : : run_info & LokiMQ : : run_info : : load ( pending_command & & pending ) {
2020-06-08 02:30:46 +02:00
if ( auto * f = std : : get_if < std : : function < void ( ) > > ( & pending . callback ) )
return load ( & pending . cat , std : : move ( pending . command ) , std : : move ( pending . remote ) , std : : move ( * f ) ) ;
assert ( pending . callback . index ( ) = = 0 ) ;
2020-04-24 02:52:39 +02:00
return load ( & pending . cat , std : : move ( pending . command ) , std : : move ( pending . conn ) , std : : move ( pending . access ) ,
2020-10-15 21:07:45 +02:00
std : : move ( pending . remote ) , std : : move ( pending . data_parts ) , var : : get < 0 > ( pending . callback ) ) ;
2020-02-28 22:54:00 +01:00
}
2020-06-05 04:32:28 +02:00
LokiMQ : : run_info & LokiMQ : : run_info : : load ( batch_job & & bj , bool reply_job , int tagged_thread ) {
2020-06-08 02:30:46 +02:00
reset ( ) ;
2020-02-06 01:21:27 +01:00
is_batch_job = true ;
2020-02-28 22:54:00 +01:00
is_reply_job = reply_job ;
2020-06-05 04:32:28 +02:00
is_tagged_thread_job = tagged_thread > 0 ;
2020-02-06 01:21:27 +01:00
batch_jobno = bj . second ;
2020-06-08 02:30:46 +02:00
to_run = bj . first ;
2020-02-06 01:21:27 +01:00
return * this ;
}
2020-02-03 03:39:26 +01:00
LokiMQ : : ~ LokiMQ ( ) {
2020-07-10 22:11:16 +02:00
if ( ! proxy_thread . joinable ( ) ) {
if ( ! tagged_workers . empty ( ) ) {
// This is a bit icky: we have tagged workers that are waiting for a signal on
// workers_socket, but the listening end of workers_socket doesn't get set up until the
// proxy thread starts (and we're getting destructed here without a proxy thread). So
// we need to start listening on it here in the destructor so that we establish a
// connection and send the QUITs to the tagged worker threads.
2020-11-10 22:47:31 +01:00
workers_socket . set ( zmq : : sockopt : : router_mandatory , true ) ;
2020-07-10 22:11:16 +02:00
workers_socket . bind ( SN_ADDR_WORKERS ) ;
2020-07-10 22:35:32 +02:00
for ( auto & [ run , busy , queue ] : tagged_workers ) {
while ( true ) {
try {
route_control ( workers_socket , run . worker_routing_id , " QUIT " ) ;
break ;
} catch ( const zmq : : error_t & ) {
std : : this_thread : : sleep_for ( 5 ms ) ;
}
}
}
2020-07-10 22:11:16 +02:00
for ( auto & [ run , busy , queue ] : tagged_workers )
run . worker_thread . join ( ) ;
}
2020-03-03 22:28:53 +01:00
return ;
2020-07-10 22:11:16 +02:00
}
2020-03-03 22:28:53 +01:00
2020-02-03 03:39:26 +01:00
LMQ_LOG ( info , " LokiMQ shutting down proxy thread " ) ;
detail : : send_control ( get_control_socket ( ) , " QUIT " ) ;
proxy_thread . join ( ) ;
LMQ_LOG ( info , " LokiMQ proxy thread has stopped " ) ;
}
2020-02-11 07:29:00 +01:00
std : : ostream & operator < < ( std : : ostream & os , LogLevel lvl ) {
os < < ( lvl = = LogLevel : : trace ? " trace " :
lvl = = LogLevel : : debug ? " debug " :
lvl = = LogLevel : : info ? " info " :
lvl = = LogLevel : : warn ? " warn " :
lvl = = LogLevel : : error ? " ERROR " :
lvl = = LogLevel : : fatal ? " FATAL " :
" unknown " ) ;
return os ;
}
2020-02-03 03:39:26 +01:00
2020-02-11 07:30:07 +01:00
std : : string make_random_string ( size_t size ) {
static thread_local std : : mt19937_64 rng { std : : random_device { } ( ) } ;
static thread_local std : : uniform_int_distribution < char > dist { std : : numeric_limits < char > : : min ( ) , std : : numeric_limits < char > : : max ( ) } ;
std : : string rando ;
rando . reserve ( size ) ;
for ( size_t i = 0 ; i < size ; i + + )
rando + = dist ( rng ) ;
return rando ;
2020-02-03 03:39:26 +01:00
}
2020-02-11 07:30:07 +01:00
} // namespace lokimq
2020-02-03 03:39:26 +01:00
// vim:sw=4:et