(cobra bots) add a utility class to factor out the common bots features (heartbeat) and move all bots to used it + convert cobra_subscribe to be a bot and add a unittest for it

This commit is contained in:
Benjamin Sergeant
2020-04-16 21:58:10 -07:00
parent 0f5d15aa11
commit a2abe861d3
20 changed files with 478 additions and 558 deletions

View File

@ -7,6 +7,7 @@ set (IXBOTS_SOURCES
ixbots/IXCobraBot.cpp
ixbots/IXCobraToSentryBot.cpp
ixbots/IXCobraToStatsdBot.cpp
ixbots/IXCobraToStdoutBot.cpp
ixbots/IXQueueManager.cpp
ixbots/IXStatsdClient.cpp
)
@ -15,6 +16,7 @@ set (IXBOTS_HEADERS
ixbots/IXCobraBot.h
ixbots/IXCobraToSentryBot.h
ixbots/IXCobraToStatsdBot.h
ixbots/IXCobraToStdoutBot.h
ixbots/IXQueueManager.h
ixbots/IXStatsdClient.h
)

View File

@ -23,6 +23,7 @@ namespace ix
const std::string& position,
bool verbose,
size_t maxQueueSize,
bool useQueue,
bool enableHeartbeat,
int runtime)
{
@ -83,16 +84,18 @@ namespace ix
std::thread t2(heartbeat);
auto sender =
[this, &queueManager, verbose, &sentCount, &stop, &throttled] {
[this, &queueManager, verbose, &sentCount, &stop, &throttled, &fatalCobraError] {
while (true)
{
Json::Value msg = queueManager.pop();
auto data = queueManager.pop();
Json::Value msg = data.first;
std::string position = data.second;
if (stop) break;
if (msg.isNull()) continue;
if (_onBotMessageCallback && _onBotMessageCallback(msg, verbose, throttled))
if (_onBotMessageCallback && _onBotMessageCallback(msg, position, verbose, throttled, fatalCobraError))
{
// That might be too noisy
if (verbose)
@ -114,16 +117,21 @@ namespace ix
std::thread t3(sender);
conn.setEventCallback([&conn,
std::string subscriptionPosition(position);
conn.setEventCallback([this,
&conn,
&channel,
&filter,
&position,
&subscriptionPosition,
&jsonWriter,
verbose,
&throttled,
&receivedCount,
&fatalCobraError,
&queueManager](const CobraEventPtr& event)
&useQueue,
&queueManager,
&sentCount](const CobraEventPtr& event)
{
if (event->type == ix::CobraEventType::Open)
{
@ -141,16 +149,20 @@ namespace ix
else if (event->type == ix::CobraEventType::Authenticated)
{
spdlog::info("Subscriber authenticated");
spdlog::info("Subscribing to {} at position {}", channel, subscriptionPosition);
spdlog::info("Using filter: {}", filter);
conn.subscribe(channel,
filter,
position,
[&jsonWriter, verbose, &throttled, &receivedCount, &queueManager](
subscriptionPosition,
[this, &jsonWriter, verbose, &throttled, &receivedCount, &queueManager, &useQueue, &subscriptionPosition, &fatalCobraError, &sentCount](
const Json::Value& msg, const std::string& position) {
if (verbose)
{
spdlog::info("Subscriber received message {} -> {}", position, jsonWriter.write(msg));
}
subscriptionPosition = position;
// If we cannot send to sentry fast enough, drop the message
if (throttled)
{
@ -158,7 +170,27 @@ namespace ix
}
++receivedCount;
queueManager.add(msg);
if (useQueue)
{
queueManager.add(msg, position);
}
else
{
if (_onBotMessageCallback && _onBotMessageCallback(msg, position, verbose, throttled, fatalCobraError))
{
// That might be too noisy
if (verbose)
{
spdlog::info("cobra bot: sending succesfull");
}
++sentCount;
}
else
{
spdlog::error("cobra bot: error sending");
}
}
});
}
else if (event->type == ix::CobraEventType::Subscribed)

View File

@ -9,15 +9,14 @@
#include <ixcobra/IXCobraConfig.h>
#include <stddef.h>
#include <json/json.h>
#include <mutex>
#include <condition_variable>
#include <queue>
#include <map>
#include <functional>
namespace ix
{
using OnBotMessageCallback = std::function<bool(const Json::Value&,
const std::string&,
const bool verbose,
std::atomic<bool>&,
std::atomic<bool>&)>;
class CobraBot
@ -31,6 +30,7 @@ namespace ix
const std::string& position,
bool verbose,
size_t maxQueueSize,
bool useQueue,
bool enableHeartbeat,
int runtime);

View File

@ -12,26 +12,26 @@
#include <ixcobra/IXCobraConnection.h>
#include <spdlog/spdlog.h>
#include <sstream>
#include <thread>
#include <vector>
namespace ix
{
int cobra_to_sentry_bot(const CobraConfig& config,
const std::string& channel,
const std::string& filter,
const std::string& position,
SentryClient& sentryClient,
bool verbose,
bool strict,
size_t maxQueueSize,
bool enableHeartbeat,
int runtime)
int64_t cobra_to_sentry_bot(const CobraConfig& config,
const std::string& channel,
const std::string& filter,
const std::string& position,
SentryClient& sentryClient,
bool verbose,
size_t maxQueueSize,
bool enableHeartbeat,
int runtime)
{
CobraBot bot;
bot.setOnBotMessageCallback([&sentryClient](const Json::Value& msg,
const std::string& /*position*/,
const bool verbose,
std::atomic<bool>& throttled) -> bool {
std::atomic<bool>& throttled,
std::atomic<bool>& /*fatalCobraError*/) -> bool {
auto ret = sentryClient.send(msg, verbose);
HttpResponsePtr response = ret.first;
@ -102,12 +102,15 @@ namespace ix
return success;
});
bool useQueue = true;
return bot.run(config,
channel,
filter,
position,
verbose,
maxQueueSize,
useQueue,
enableHeartbeat,
runtime);
}

View File

@ -8,17 +8,17 @@
#include <ixcobra/IXCobraConfig.h>
#include <ixsentry/IXSentryClient.h>
#include <string>
#include <cstdint>
namespace ix
{
int cobra_to_sentry_bot(const CobraConfig& config,
const std::string& channel,
const std::string& filter,
const std::string& position,
SentryClient& sentryClient,
bool verbose,
bool strict,
size_t maxQueueSize,
bool enableHeartbeat,
int runtime);
int64_t cobra_to_sentry_bot(const CobraConfig& config,
const std::string& channel,
const std::string& filter,
const std::string& position,
SentryClient& sentryClient,
bool verbose,
size_t maxQueueSize,
bool enableHeartbeat,
int runtime);
} // namespace ix

View File

@ -7,14 +7,12 @@
#include "IXCobraToStatsdBot.h"
#include "IXQueueManager.h"
#include "IXStatsdClient.h"
#include "IXCobraBot.h"
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <ixcobra/IXCobraConnection.h>
#include <spdlog/spdlog.h>
#include <sstream>
#include <thread>
#include <vector>
namespace ix
@ -56,18 +54,18 @@ namespace ix
return val;
}
int cobra_to_statsd_bot(const ix::CobraConfig& config,
const std::string& channel,
const std::string& filter,
const std::string& position,
StatsdClient& statsdClient,
const std::string& fields,
const std::string& gauge,
const std::string& timer,
bool verbose,
size_t maxQueueSize,
bool enableHeartbeat,
int runtime)
int64_t cobra_to_statsd_bot(const ix::CobraConfig& config,
const std::string& channel,
const std::string& filter,
const std::string& position,
StatsdClient& statsdClient,
const std::string& fields,
const std::string& gauge,
const std::string& timer,
bool verbose,
size_t maxQueueSize,
bool enableHeartbeat,
int runtime)
{
ix::CobraConnection conn;
conn.configure(config);
@ -75,242 +73,85 @@ namespace ix
auto tokens = parseFields(fields);
Json::FastWriter jsonWriter;
std::atomic<uint64_t> sentCount(0);
std::atomic<uint64_t> receivedCount(0);
std::atomic<bool> stop(false);
std::atomic<bool> fatalCobraError(false);
QueueManager queueManager(maxQueueSize);
auto progress = [&sentCount, &receivedCount, &stop] {
while (!stop)
CobraBot bot;
bot.setOnBotMessageCallback([&statsdClient, &tokens, &gauge, &timer](const Json::Value& msg,
const std::string& /*position*/,
const bool verbose,
std::atomic<bool>& /*throttled*/,
std::atomic<bool>& fatalCobraError) -> bool {
std::string id;
for (auto&& attr : tokens)
{
spdlog::info("messages received {} sent {}", receivedCount, sentCount);
auto duration = std::chrono::seconds(1);
std::this_thread::sleep_for(duration);
id += ".";
auto val = extractAttr(attr, msg);
id += val.asString();
}
spdlog::info("timer thread done");
};
std::thread t1(progress);
auto heartbeat = [&sentCount, &receivedCount, &stop, &enableHeartbeat] {
std::string state("na");
if (!enableHeartbeat) return;
while (!stop)
if (gauge.empty() && timer.empty())
{
std::stringstream ss;
ss << "messages received " << receivedCount;
ss << "messages sent " << sentCount;
std::string currentState = ss.str();
if (currentState == state)
{
spdlog::error("no messages received or sent for 1 minute, exiting");
exit(1);
}
state = currentState;
auto duration = std::chrono::minutes(1);
std::this_thread::sleep_for(duration);
statsdClient.count(id, 1);
}
spdlog::info("heartbeat thread done");
};
std::thread t2(heartbeat);
auto statsdSender = [&statsdClient, &queueManager, &sentCount, &tokens, &stop, &gauge, &timer, &fatalCobraError, &verbose] {
while (true)
else
{
Json::Value msg = queueManager.pop();
std::string attrName = (!gauge.empty()) ? gauge : timer;
auto val = extractAttr(attrName, msg);
size_t x;
if (stop) return;
if (msg.isNull()) continue;
std::string id;
for (auto&& attr : tokens)
if (val.isInt())
{
id += ".";
auto val = extractAttr(attr, msg);
id += val.asString();
x = (size_t) val.asInt();
}
if (gauge.empty() && timer.empty())
else if (val.isInt64())
{
statsdClient.count(id, 1);
x = (size_t) val.asInt64();
}
else if (val.isUInt())
{
x = (size_t) val.asUInt();
}
else if (val.isUInt64())
{
x = (size_t) val.asUInt64();
}
else if (val.isDouble())
{
x = (size_t) val.asUInt64();
}
else
{
std::string attrName = (!gauge.empty()) ? gauge : timer;
auto val = extractAttr(attrName, msg);
size_t x;
if (val.isInt())
{
x = (size_t) val.asInt();
}
else if (val.isInt64())
{
x = (size_t) val.asInt64();
}
else if (val.isUInt())
{
x = (size_t) val.asUInt();
}
else if (val.isUInt64())
{
x = (size_t) val.asUInt64();
}
else if (val.isDouble())
{
x = (size_t) val.asUInt64();
}
else
{
spdlog::error("Gauge {} is not a numberic type", gauge);
fatalCobraError = true;
break;
}
if (verbose)
{
spdlog::info("{} - {} -> {}", id, attrName, x);
}
if (!gauge.empty())
{
statsdClient.gauge(id, x);
}
else
{
statsdClient.timing(id, x);
}
}
sentCount += 1;
}
};
std::thread t3(statsdSender);
conn.setEventCallback(
[&conn, &channel, &filter, &position, &jsonWriter, verbose, &queueManager, &receivedCount, &fatalCobraError](const CobraEventPtr& event)
{
if (event->type == ix::CobraEventType::Open)
{
spdlog::info("Subscriber connected");
for (auto&& it : event->headers)
{
spdlog::info("{}: {}", it.first, it.second);
}
}
else if (event->type == ix::CobraEventType::Closed)
{
spdlog::info("Subscriber closed: {}", event->errMsg);
}
else if (event->type == ix::CobraEventType::Authenticated)
{
spdlog::info("Subscriber authenticated");
conn.subscribe(channel,
filter,
position,
[&jsonWriter, &queueManager, verbose, &receivedCount](
const Json::Value& msg, const std::string& position) {
if (verbose)
{
spdlog::info("Subscriber received message {} -> {}", position, jsonWriter.write(msg));
}
receivedCount++;
++receivedCount;
queueManager.add(msg);
});
}
else if (event->type == ix::CobraEventType::Subscribed)
{
spdlog::info("Subscriber: subscribed to channel {}", event->subscriptionId);
}
else if (event->type == ix::CobraEventType::UnSubscribed)
{
spdlog::info("Subscriber: unsubscribed from channel {}", event->subscriptionId);
}
else if (event->type == ix::CobraEventType::Error)
{
spdlog::error("Subscriber: error {}", event->errMsg);
}
else if (event->type == ix::CobraEventType::Published)
{
spdlog::error("Published message hacked: {}", event->msgId);
}
else if (event->type == ix::CobraEventType::Pong)
{
spdlog::info("Received websocket pong");
}
else if (event->type == ix::CobraEventType::HandshakeError)
{
spdlog::error("Subscriber: Handshake error: {}", event->errMsg);
spdlog::error("Gauge {} is not a numeric type", gauge);
fatalCobraError = true;
return false;
}
else if (event->type == ix::CobraEventType::AuthenticationError)
if (verbose)
{
spdlog::error("Subscriber: Authentication error: {}", event->errMsg);
fatalCobraError = true;
spdlog::info("{} - {} -> {}", id, attrName, x);
}
else if (event->type == ix::CobraEventType::SubscriptionError)
if (!gauge.empty())
{
spdlog::error("Subscriber: Subscription error: {}", event->errMsg);
fatalCobraError = true;
statsdClient.gauge(id, x);
}
else
{
statsdClient.timing(id, x);
}
});
// Run forever
if (runtime == -1)
{
while (true)
{
auto duration = std::chrono::seconds(1);
std::this_thread::sleep_for(duration);
if (fatalCobraError) break;
}
}
// Run for a duration, used by unittesting now
else
{
for (int i = 0 ; i < runtime; ++i)
{
auto duration = std::chrono::seconds(1);
std::this_thread::sleep_for(duration);
if (fatalCobraError) break;
}
}
return true;
});
//
// Cleanup.
// join all the bg threads and stop them.
//
conn.disconnect();
stop = true;
bool useQueue = true;
// progress thread
t1.join();
// heartbeat thread
if (t2.joinable()) t2.join();
// statsd sender thread
t3.join();
return fatalCobraError ? -1 : (int) sentCount;
return bot.run(config,
channel,
filter,
position,
verbose,
maxQueueSize,
useQueue,
enableHeartbeat,
runtime);
}
} // namespace ix

View File

@ -9,19 +9,20 @@
#include <ixbots/IXStatsdClient.h>
#include <string>
#include <stddef.h>
#include <cstdint>
namespace ix
{
int cobra_to_statsd_bot(const ix::CobraConfig& config,
const std::string& channel,
const std::string& filter,
const std::string& position,
StatsdClient& statsdClient,
const std::string& fields,
const std::string& gauge,
const std::string& timer,
bool verbose,
size_t maxQueueSize,
bool enableHeartbeat,
int runtime);
int64_t cobra_to_statsd_bot(const ix::CobraConfig& config,
const std::string& channel,
const std::string& filter,
const std::string& position,
StatsdClient& statsdClient,
const std::string& fields,
const std::string& gauge,
const std::string& timer,
bool verbose,
size_t maxQueueSize,
bool enableHeartbeat,
int runtime);
} // namespace ix

View File

@ -0,0 +1,106 @@
/*
* IXCobraToStdoutBot.cpp
* Author: Benjamin Sergeant
* Copyright (c) 2019 Machine Zone, Inc. All rights reserved.
*/
#include "IXCobraToStdoutBot.h"
#include "IXCobraBot.h"
#include "IXQueueManager.h"
#include <chrono>
#include <spdlog/spdlog.h>
#include <sstream>
#include <iostream>
namespace ix
{
using StreamWriterPtr = std::unique_ptr<Json::StreamWriter>;
StreamWriterPtr makeStreamWriter()
{
Json::StreamWriterBuilder builder;
builder["commentStyle"] = "None";
builder["indentation"] = ""; // will make the JSON object compact
std::unique_ptr<Json::StreamWriter> jsonWriter(builder.newStreamWriter());
return jsonWriter;
}
std::string timeSinceEpoch()
{
std::chrono::system_clock::time_point tp = std::chrono::system_clock::now();
std::chrono::system_clock::duration dtn = tp.time_since_epoch();
std::stringstream ss;
ss << dtn.count() * std::chrono::system_clock::period::num /
std::chrono::system_clock::period::den;
return ss.str();
}
void writeToStdout(bool fluentd,
const StreamWriterPtr& jsonWriter,
const Json::Value& msg,
const std::string& position)
{
Json::Value enveloppe;
if (fluentd)
{
enveloppe["producer"] = "cobra";
enveloppe["consumer"] = "fluentd";
Json::Value nestedMessage(msg);
nestedMessage["position"] = position;
nestedMessage["created_at"] = timeSinceEpoch();
enveloppe["message"] = nestedMessage;
jsonWriter->write(enveloppe, &std::cout);
std::cout << std::endl; // add lf and flush
}
else
{
enveloppe = msg;
std::cout << position << " ";
jsonWriter->write(enveloppe, &std::cout);
std::cout << std::endl;
}
}
int64_t cobra_to_stdout_bot(const CobraConfig& config,
const std::string& channel,
const std::string& filter,
const std::string& position,
bool fluentd,
bool quiet,
bool verbose,
size_t maxQueueSize,
bool enableHeartbeat,
int runtime)
{
CobraBot bot;
auto jsonWriter = makeStreamWriter();
bot.setOnBotMessageCallback([&fluentd, &quiet, &jsonWriter](const Json::Value& msg,
const std::string& position,
const bool /*verbose*/,
std::atomic<bool>& /*throttled*/,
std::atomic<bool>& /*fatalCobraError*/) -> bool {
if (!quiet)
{
writeToStdout(fluentd, jsonWriter, msg, position);
}
return true;
});
bool useQueue = false;
return bot.run(config,
channel,
filter,
position,
verbose,
maxQueueSize,
useQueue,
enableHeartbeat,
runtime);
}
} // namespace ix

View File

@ -0,0 +1,25 @@
/*
* IXCobraToStdoutBot.h
* Author: Benjamin Sergeant
* Copyright (c) 2019-2020 Machine Zone, Inc. All rights reserved.
*/
#pragma once
#include <ixcobra/IXCobraConfig.h>
#include <string>
#include <stddef.h>
#include <cstdint>
namespace ix
{
int64_t cobra_to_stdout_bot(const ix::CobraConfig& config,
const std::string& channel,
const std::string& filter,
const std::string& position,
bool fluentd,
bool quiet,
bool verbose,
size_t maxQueueSize,
bool enableHeartbeat,
int runtime);
} // namespace ix

View File

@ -10,14 +10,14 @@
namespace ix
{
Json::Value QueueManager::pop()
std::pair<Json::Value, std::string> QueueManager::pop()
{
std::unique_lock<std::mutex> lock(_mutex);
if (_queues.empty())
{
Json::Value val;
return val;
return std::make_pair(val, std::string());
}
std::vector<std::string> games;
@ -35,7 +35,7 @@ namespace ix
if (_queues[game].empty())
{
Json::Value val;
return val;
return std::make_pair(val, std::string());
}
auto msg = _queues[game].front();
@ -43,7 +43,7 @@ namespace ix
return msg;
}
void QueueManager::add(Json::Value msg)
void QueueManager::add(const Json::Value& msg, const std::string& position)
{
std::unique_lock<std::mutex> lock(_mutex);
@ -59,7 +59,7 @@ namespace ix
// in queuing too many events.
if (_queues[game].size() < _maxQueueSize)
{
_queues[game].push(msg);
_queues[game].push(std::make_pair(msg, position));
_condition.notify_one();
}
}

View File

@ -23,11 +23,11 @@ namespace ix
{
}
Json::Value pop();
void add(Json::Value msg);
std::pair<Json::Value, std::string> pop();
void add(const Json::Value& msg, const std::string& position);
private:
std::map<std::string, std::queue<Json::Value>> _queues;
std::map<std::string, std::queue<std::pair<Json::Value, std::string>>> _queues;
std::mutex _mutex;
std::condition_variable _condition;
size_t _maxQueueSize;