├── .gitignore
├── include
    ├── python
    │   ├── api.h
    │   ├── plugin.h
    │   ├── helpers.h
    │   ├── lag_tracker_handler.h
    │   └── handler.h
    ├── utils
    │   ├── utils.h
    │   ├── thread_pool.h
    │   ├── async_observer.h
    │   ├── task_scheduler.h
    │   └── observer.h
    ├── detail
    │   ├── logging.h
    │   ├── memory.h
    │   └── endianness.h
    ├── exceptions.h
    ├── plugin_base.h
    ├── consumer_pool.h
    ├── application.h
    ├── consumer_offset.h
    ├── consumer_offset_reader.h
    ├── topic_offset_reader.h
    └── offset_store.h
├── CMakeLists.txt
├── src
    ├── plugin_base.cpp
    ├── detail
    │   └── logging.cpp
    ├── utils
    │   ├── utils.cpp
    │   ├── thread_pool.cpp
    │   └── task_scheduler.cpp
    ├── CMakeLists.txt
    ├── consumer_pool.cpp
    ├── consumer_offset.cpp
    ├── application.cpp
    ├── python
    │   ├── helpers.cpp
    │   ├── plugin.cpp
    │   ├── lag_tracker_handler.cpp
    │   ├── handler.cpp
    │   └── api.cpp
    ├── main.cpp
    ├── consumer_offset_reader.cpp
    ├── offset_store.cpp
    └── topic_offset_reader.cpp
├── plugins
    ├── lag_exporter.py
    └── logger.py
├── executables
    ├── CMakeLists.txt
    ├── topic_offsets.cpp
    └── consumer_offsets.cpp
└── LICENSE


/.gitignore:
--------------------------------------------------------------------------------
1 | build
2 | *.pyc
3 | 


--------------------------------------------------------------------------------
/include/python/api.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | namespace pirulo {
 4 | namespace api {
 5 | 
 6 | void register_module();
 7 | void register_types();
 8 | 
 9 | } // api
10 | } // pirulo
11 | 


--------------------------------------------------------------------------------
/CMakeLists.txt:
--------------------------------------------------------------------------------
1 | cmake_minimum_required(VERSION 2.8.1)
2 | project(pirulo)
3 | 
4 | set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall")
5 | add_subdirectory(src)
6 | add_subdirectory(executables)
7 | 


--------------------------------------------------------------------------------
/include/utils/utils.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | #include <string>
 4 | 
 5 | namespace pirulo {
 6 | namespace utils {
 7 | 
 8 | std::string generate_group_id();
 9 | 
10 | } // utils
11 | } // pirulo
12 | 


--------------------------------------------------------------------------------
/src/plugin_base.cpp:
--------------------------------------------------------------------------------
 1 | #include "plugin_base.h"
 2 | 
 3 | using std::move;
 4 | 
 5 | namespace pirulo {
 6 | 
 7 | void PluginBase::launch(StorePtr store) {
 8 |     store_ = move(store);
 9 |     initialize();
10 | }
11 | 
12 | PluginBase::StorePtr PluginBase::get_store() const {
13 |     return store_;
14 | }
15 | 
16 | } // pirulo
17 | 


--------------------------------------------------------------------------------
/plugins/lag_exporter.py:
--------------------------------------------------------------------------------
 1 | from pirulo import LagTrackerHandler
 2 | 
 3 | class Plugin(LagTrackerHandler):
 4 |     def __init__(self):
 5 |         LagTrackerHandler.__init__(self)
 6 | 
 7 |     def handle_lag_update(self, topic, partition, group_id, lag):
 8 |         print 'Consumer {0} has {1} lag on {2}/{3}'.format(group_id, lag, topic, partition)
 9 | 
10 | def create_plugin():
11 |     return Plugin()
12 | 


--------------------------------------------------------------------------------
/include/detail/logging.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | #include <string>
 4 | #include <log4cxx/logger.h>
 5 | 
 6 | #define PIRULO_CREATE_LOGGER(name) static const auto logger = log4cxx::Logger::getLogger(name)
 7 | 
 8 | namespace pirulo {
 9 | namespace logging {
10 | 
11 | void register_console_logger(const std::string& name = "",
12 |                              const std::string& log_level = "DEBUG");
13 | 
14 | } // logging
15 | } // pirulo
16 | 


--------------------------------------------------------------------------------
/include/exceptions.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | #include <stdexcept>
 4 | 
 5 | namespace pirulo {
 6 | 
 7 | class Exception : public std::runtime_error {
 8 | public:
 9 |     using std::runtime_error::runtime_error;
10 | };
11 | 
12 | class ParseException : public Exception {
13 | public:
14 |     using Exception::Exception;
15 | 
16 |     ParseException() : Exception("Parsing record failed") {
17 | 
18 |     }
19 | };
20 | 
21 | } // pirulo
22 | 


--------------------------------------------------------------------------------
/include/plugin_base.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | #include <memory>
 4 | #include "offset_store.h"
 5 | 
 6 | namespace pirulo {
 7 | 
 8 | class PluginBase {
 9 | public:
10 |     using StorePtr = std::shared_ptr<OffsetStore>;
11 | 
12 |     virtual ~PluginBase() = default;
13 | 
14 |     void launch(StorePtr store);
15 | protected:
16 |     StorePtr get_store() const;
17 | private:
18 |     virtual void initialize() = 0;
19 | 
20 |     StorePtr store_;
21 | };
22 | 
23 | } // pirulo
24 | 


--------------------------------------------------------------------------------
/include/python/plugin.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | #include <string>
 4 | #include <boost/python/object.hpp>
 5 | #include "plugin_base.h"
 6 | 
 7 | namespace pirulo {
 8 | namespace api {
 9 | 
10 | class PythonPlugin : public PluginBase {
11 | public:
12 |     PythonPlugin(const std::string& modules_path, const std::string& file_path);
13 |     ~PythonPlugin();
14 | private:
15 |     void initialize();
16 | 
17 |     boost::python::object plugin_;
18 | };
19 | 
20 | } // api
21 | } // pirulo
22 | 


--------------------------------------------------------------------------------
/executables/CMakeLists.txt:
--------------------------------------------------------------------------------
 1 | set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin/utils)
 2 | 
 3 | add_custom_target(executables)
 4 | 
 5 | macro(create_executable executable_name)
 6 |     add_executable(${executable_name} EXCLUDE_FROM_ALL "${executable_name}.cpp")
 7 |     add_dependencies(executables ${executable_name})
 8 |     target_link_libraries(${executable_name} pirulo-core)
 9 | endmacro()
10 | 
11 | include_directories(${PROJECT_SOURCE_DIR}/include)
12 | create_executable(consumer_offsets)
13 | create_executable(topic_offsets)
14 | 


--------------------------------------------------------------------------------
/src/detail/logging.cpp:
--------------------------------------------------------------------------------
 1 | #include <log4cxx/patternlayout.h>
 2 | #include <log4cxx/consoleappender.h>
 3 | #include "detail/logging.h"
 4 | 
 5 | using std::string;
 6 | 
 7 | using log4cxx::PatternLayout;
 8 | using log4cxx::ConsoleAppender;
 9 | using log4cxx::Level;
10 | using log4cxx::Logger;
11 | 
12 | namespace pirulo {
13 | namespace logging {
14 | 
15 | void register_console_logger(const string& name, const string& log_level) {
16 |     auto layout = new PatternLayout("%d{yyyy-MM-dd HH:mm:ss.SSS}{GMT} [%c{2}] - %m%n");
17 |     auto appender = new ConsoleAppender(layout);
18 | 
19 |     auto logger = Logger::getRootLogger();
20 |     logger->setLevel(Level::toLevel(log_level));
21 |     logger->addAppender(appender);
22 | }
23 | 
24 | } // logging
25 | } // pirulo
26 | 


--------------------------------------------------------------------------------
/include/consumer_pool.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | #include <vector>
 4 | #include <mutex>
 5 | #include <condition_variable>
 6 | #include <functional>
 7 | #include <queue>
 8 | #include <cppkafka/consumer.h>
 9 | 
10 | namespace pirulo {
11 | 
12 | class ConsumerPool {
13 | public:
14 |     using ConsumerCallback = std::function<void(cppkafka::Consumer&)>;
15 | 
16 |     ConsumerPool(size_t consumer_count, cppkafka::Configuration config);
17 | 
18 |     void acquire_consumer(const ConsumerCallback& callback);
19 | private:
20 |     using ConsumerContainer = std::deque<cppkafka::Consumer>;
21 |     ConsumerContainer consumers_;
22 |     std::queue<ConsumerContainer::reverse_iterator> available_consumers_;
23 |     std::mutex consumers_mutex_;
24 |     std::condition_variable consumers_condition_;
25 | };
26 | 
27 | } // pirulo
28 | 


--------------------------------------------------------------------------------
/include/application.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | #include <memory>
 4 | #include <vector>
 5 | #include "topic_offset_reader.h"
 6 | #include "consumer_offset_reader.h"
 7 | #include "plugin_base.h"
 8 | 
 9 | namespace pirulo {
10 | 
11 | class Application {
12 | public:
13 |     using TopicOffsetReaderPtr = std::shared_ptr<TopicOffsetReader>;
14 |     using ConsumerOffsetReaderPtr = std::shared_ptr<ConsumerOffsetReader>;
15 |     using PluginPtr = std::unique_ptr<PluginBase>;
16 | 
17 |     Application(TopicOffsetReaderPtr topic_reader, ConsumerOffsetReaderPtr consumer_reader);
18 | 
19 |     void run();
20 |     void stop();
21 | 
22 |     void add_plugin(PluginPtr plugin);
23 | private:
24 |     void process();
25 | 
26 |     std::vector<PluginPtr> plugins_;
27 |     TopicOffsetReaderPtr topic_reader_;
28 |     ConsumerOffsetReaderPtr consumer_reader_;
29 | };
30 | 
31 | } // pirulo
32 | 


--------------------------------------------------------------------------------
/include/consumer_offset.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | #include <string>
 4 | #include <memory>
 5 | #include <cppkafka/topic_partition.h>
 6 | 
 7 | namespace pirulo {
 8 | 
 9 | class ConsumerOffset {
10 | public:
11 |     ConsumerOffset(std::string group_id, std::string topic, int partition,
12 |                    uint64_t offset);
13 |     const std::string& get_group_id() const;
14 |     const cppkafka::TopicPartition& get_topic_partition() const;
15 | private:
16 |     std::string group_id_;
17 |     cppkafka::TopicPartition topic_partition_;
18 | };
19 | 
20 | bool operator==(const ConsumerOffset& lhs, const ConsumerOffset& rhs);
21 | bool operator!=(const ConsumerOffset& lhs, const ConsumerOffset& rhs);
22 | 
23 | } // pirulo
24 | 
25 | // Specialize std::hash for ConsumerOffset
26 | namespace std {
27 | 
28 | template <>
29 | struct hash<pirulo::ConsumerOffset> {
30 |     size_t operator()(const pirulo::ConsumerOffset& consumer_offset) const;
31 | };
32 | 
33 | } // std
34 | 


--------------------------------------------------------------------------------
/src/utils/utils.cpp:
--------------------------------------------------------------------------------
 1 | #include <random>
 2 | #include <array>
 3 | #include "utils/utils.h"
 4 | 
 5 | using std::string;
 6 | using std::array;
 7 | using std::random_device;
 8 | using std::mt19937;
 9 | using std::uniform_int_distribution;
10 | 
11 | namespace pirulo {
12 | namespace utils {
13 | 
14 | static const size_t GROUP_ID_LENGTH = 16;
15 | 
16 | string generate_group_id() {
17 |     array<char, 16> alphanums;
18 |     for (size_t i = 0; i <= 9; ++i) {
19 |         alphanums[i] = '0' + i;
20 |     }
21 |     for (size_t c = 'a'; c <= 'f'; ++c) {
22 |         alphanums[10 + c - 'a'] = c;
23 |     }
24 |     const auto seed = random_device{}();
25 |     mt19937 generator(seed);
26 |     uniform_int_distribution<size_t> distribution(0, alphanums.size() - 1);
27 | 
28 |     string output = "pirulo-";
29 |     for (size_t i = 0; i < GROUP_ID_LENGTH; ++i) {
30 |         output.push_back(alphanums[distribution(generator)]);
31 |     }
32 |     return output;
33 | }
34 | 
35 | } // utils
36 | } // pirulo
37 | 
38 | 


--------------------------------------------------------------------------------
/src/CMakeLists.txt:
--------------------------------------------------------------------------------
 1 | set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin)
 2 | 
 3 | set(SOURCES
 4 |     consumer_offset.cpp
 5 |     offset_store.cpp
 6 |     consumer_offset_reader.cpp
 7 |     topic_offset_reader.cpp
 8 |     plugin_base.cpp
 9 |     consumer_pool.cpp
10 |     application.cpp
11 |     
12 |     utils/thread_pool.cpp
13 |     utils/task_scheduler.cpp
14 |     utils/utils.cpp
15 | 
16 |     detail/logging.cpp
17 | 
18 |     python/plugin.cpp
19 |     python/helpers.cpp
20 |     python/handler.cpp
21 |     python/lag_tracker_handler.cpp
22 |     python/api.cpp
23 | )
24 | 
25 | include_directories(${PROJECT_SOURCE_DIR}/include)
26 | # TODO: clean this up
27 | include_directories(SYSTEM /usr/include/python2.7/)
28 | 
29 | add_library(pirulo-core ${SOURCES})
30 | target_link_libraries(pirulo-core log4cxx cppkafka rdkafka boost_program_options boost_python
31 |                                   python2.7 pthread)
32 | 
33 | add_executable(pirulo main.cpp)
34 | target_link_libraries(pirulo pirulo-core )
35 | 


--------------------------------------------------------------------------------
/include/python/helpers.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | #include <string>
 4 | #include <log4cxx/logger.h>
 5 | #include <boost/python/errors.hpp>
 6 | #include <pystate.h>
 7 | 
 8 | namespace pirulo {
 9 | namespace api {
10 | namespace helpers {
11 | 
12 | class GILAcquirer {
13 | public:
14 |     GILAcquirer();
15 |     ~GILAcquirer();
16 | 
17 |     GILAcquirer(const GILAcquirer&) = delete;
18 |     GILAcquirer& operator=(const GILAcquirer&) = delete;
19 | private:
20 |     PyGILState_STATE gstate_;
21 | };
22 | 
23 | std::string format_python_exception();
24 | void initialize_python();
25 | 
26 | template <typename Functor>
27 | void safe_exec(const log4cxx::LoggerPtr& logger, const Functor& python_code) {
28 |     GILAcquirer _;
29 |     try {
30 |         python_code();
31 |     }
32 |     catch (const boost::python::error_already_set& ex) {
33 |         LOG4CXX_ERROR(logger, "Error executing python callback: "
34 |                       << format_python_exception());
35 |     }
36 | }
37 | 
38 | } // helpers
39 | } // api
40 | } // pirulo
41 | 


--------------------------------------------------------------------------------
/src/consumer_pool.cpp:
--------------------------------------------------------------------------------
 1 | #include "consumer_pool.h"
 2 | 
 3 | using std::unique_lock;
 4 | using std::mutex;
 5 | 
 6 | using cppkafka::Configuration;
 7 | using cppkafka::Consumer;
 8 | 
 9 | namespace pirulo {
10 | 
11 | ConsumerPool::ConsumerPool(size_t consumer_count, Configuration config) {
12 |     for (size_t i = 0; i < consumer_count; ++i) {
13 |         consumers_.emplace_back(config);
14 |         available_consumers_.push(consumers_.rbegin());
15 |     }
16 | }
17 | 
18 | void ConsumerPool::acquire_consumer(const ConsumerCallback& callback) {
19 |     unique_lock<mutex> lock(consumers_mutex_);
20 |     while (available_consumers_.empty()) {
21 |         consumers_condition_.wait(lock);
22 |     }
23 |     auto iter = available_consumers_.front();
24 |     available_consumers_.pop();
25 |     Consumer& consumer = *iter;
26 | 
27 |     // Execute outside of the critical section
28 |     lock.unlock();
29 |     callback(consumer);
30 |     lock.lock();
31 | 
32 |     available_consumers_.push(iter);
33 |     consumers_condition_.notify_one();
34 | }
35 | 
36 | } // pirulo
37 | 


--------------------------------------------------------------------------------
/include/utils/thread_pool.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | #include <vector>
 4 | #include <thread>
 5 | #include <functional>
 6 | #include <queue>
 7 | #include <mutex>
 8 | #include <condition_variable>
 9 | #include <atomic>
10 | 
11 | namespace pirulo {
12 | 
13 | class ThreadPool {
14 | public:
15 |     using Task = std::function<void()>;
16 | 
17 |     ThreadPool(size_t thread_count);
18 |     ThreadPool(size_t thread_count, size_t maximum_tasks);
19 |     ThreadPool(const ThreadPool&) = delete;
20 |     ThreadPool& operator=(const ThreadPool&) = delete;
21 |     ~ThreadPool();
22 | 
23 |     // Returns true iff the task was successfully added. Adding a task will only fail
24 |     // iff a maximum task limit has been set and the limit has been reached
25 |     bool add_task(Task task);
26 |     void stop();
27 |     void wait_for_tasks();
28 | private:
29 |     void process();
30 | 
31 |     std::vector<std::thread> threads_;
32 |     std::queue<Task> tasks_;
33 |     std::mutex tasks_mutex_;
34 |     std::condition_variable tasks_condition_;
35 |     std::condition_variable no_tasks_condition_;
36 |     const size_t maximum_tasks_;
37 |     std::atomic<bool> running_{true};
38 | };
39 | 
40 | } // pirulo
41 | 


--------------------------------------------------------------------------------
/include/consumer_offset_reader.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | #include <memory>
 4 | #include <set>
 5 | #include <cppkafka/consumer.h>
 6 | #include <cppkafka/utils/consumer_dispatcher.h>
 7 | #include "offset_store.h"
 8 | #include "utils/observer.h"
 9 | 
10 | namespace pirulo {
11 | 
12 | class ConsumerOffsetReader {
13 | public:
14 |     using StorePtr = std::shared_ptr<OffsetStore>;
15 |     using EofCallback = std::function<void()>;
16 |     using TopicCommitCallback = std::function<void(const std::string&, int)>;
17 | 
18 |     ConsumerOffsetReader(StorePtr store, std::chrono::milliseconds consumer_offset_cool_down,
19 |                          cppkafka::Configuration config);
20 | 
21 |     void run(const EofCallback& callback);
22 |     void stop();
23 | 
24 |     void watch_commits(const std::string& topic, int partition, TopicCommitCallback callback);
25 | 
26 |     StorePtr get_store() const;
27 | private:
28 |     void handle_message(const cppkafka::Message& msg);
29 | 
30 |     StorePtr store_;
31 |     cppkafka::Consumer consumer_;
32 |     cppkafka::ConsumerDispatcher dispatcher_{consumer_};
33 |     Observer<cppkafka::TopicPartition> observer_;
34 |     std::set<int> pending_partitions_;
35 |     bool notifications_enabled_{false};
36 | };
37 | 
38 | } // pirulo
39 | 


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
 1 | Copyright (c) 2017, Matias Fontanini
 2 | All rights reserved.
 3 | 
 4 | Redistribution and use in source and binary forms, with or without
 5 | modification, are permitted provided that the following conditions are
 6 | met:
 7 | 
 8 | * Redistributions of source code must retain the above copyright
 9 |   notice, this list of conditions and the following disclaimer.
10 | * Redistributions in binary form must reproduce the above
11 |   copyright notice, this list of conditions and the following disclaimer
12 |   in the documentation and/or other materials provided with the
13 |   distribution.
14 | 
15 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
16 | "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
17 | LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
18 | A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
19 | OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
20 | SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
21 | LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
22 | DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
23 | THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
24 | (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
25 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.


--------------------------------------------------------------------------------
/include/python/lag_tracker_handler.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | #include <tuple>
 4 | #include <map>
 5 | #include <unordered_map>
 6 | #include "python/handler.h"
 7 | 
 8 | namespace pirulo {
 9 | namespace api {
10 | 
11 | class LagTrackerHandler : public Handler {
12 | public:
13 |     using Handler::Handler;
14 | protected:
15 |     virtual void handle_lag_update(const std::string& topic, int partition,
16 |                                    const std::string& group_id, uint64_t consumer_lag) { }
17 | 
18 |     void handle_initialize() override;
19 |     void handle_new_consumer(const std::string& group_id) override;
20 |     void handle_new_topic(const std::string& topic) override;
21 |     void handle_consumer_commit(const std::string& group_id, const std::string& topic,
22 |                                 int partition, int64_t offset) override;
23 |     void handle_topic_message(const std::string& topic, int partition,
24 |                               int64_t offset) override;
25 | private:
26 |     struct TopicPartitionInfo {
27 |         int64_t offset{-1};
28 |         std::unordered_map<std::string, int64_t> consumer_offsets;
29 |     };
30 |     using TopicPartitionId = std::tuple<std::string, int>;
31 |     using TopicPartitionInfoMap = std::map<TopicPartitionId, TopicPartitionInfo>;
32 | 
33 |     TopicPartitionInfoMap topic_partition_info_;
34 | };
35 | 
36 | } // api
37 | } // pirulo
38 | 


--------------------------------------------------------------------------------
/src/consumer_offset.cpp:
--------------------------------------------------------------------------------
 1 | #include <boost/functional/hash.hpp>
 2 | #include "consumer_offset.h"
 3 | 
 4 | using std::string;
 5 | 
 6 | using cppkafka::TopicPartition;
 7 | 
 8 | namespace pirulo {
 9 | 
10 | ConsumerOffset::ConsumerOffset(string group_id, string topic,
11 |                                int partition, uint64_t offset)
12 | : group_id_(move(group_id)), topic_partition_(move(topic), partition, offset) {
13 | 
14 | }
15 | 
16 | const string& ConsumerOffset::get_group_id() const {
17 |     return group_id_;
18 | }
19 | 
20 | const TopicPartition& ConsumerOffset::get_topic_partition() const {
21 |     return topic_partition_;
22 | }
23 | 
24 | bool operator==(const ConsumerOffset& lhs, const ConsumerOffset& rhs) {
25 |     return lhs.get_topic_partition() == rhs.get_topic_partition();
26 | }
27 | 
28 | bool operator!=(const ConsumerOffset& lhs, const ConsumerOffset& rhs) {
29 |     return !(lhs == rhs);
30 | }
31 | 
32 | } // pirulo
33 | 
34 | using CO = pirulo::ConsumerOffset;
35 | 
36 | namespace std {
37 | 
38 | size_t hash<CO>::operator()(const CO& consumer_offset) const {
39 |     size_t output = 0;
40 |     boost::hash_combine(output, consumer_offset.get_group_id());
41 |     boost::hash_combine(output, consumer_offset.get_topic_partition().get_topic());
42 |     boost::hash_combine(output, consumer_offset.get_topic_partition().get_partition());
43 |     return output;
44 | }
45 | 
46 | } // std
47 | 


--------------------------------------------------------------------------------
/include/python/handler.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | #include "offset_store.h"
 4 | 
 5 | namespace pirulo {
 6 | namespace api {
 7 | 
 8 | class Handler {
 9 | public:
10 |     virtual ~Handler() = default;
11 | 
12 |     void initialize(const std::shared_ptr<OffsetStore>& store);
13 |     void subscribe_to_consumers();
14 |     void subscribe_to_consumer_commits();
15 |     void subscribe_to_topics();
16 |     void subscribe_to_topic_message();
17 |     const std::shared_ptr<OffsetStore>& get_offset_store() const;
18 | protected:
19 |     virtual void handle_initialize(); 
20 |     virtual void handle_new_consumer(const std::string& group_id);
21 |     virtual void handle_new_topic(const std::string& topic);
22 |     virtual void handle_consumer_commit(const std::string& group_id, const std::string& topic,
23 |                                         int partition, int64_t offset);
24 |     virtual void handle_topic_message(const std::string& topic, int partition, int64_t offset);
25 | private:
26 |     void on_new_consumer(const std::string& group_id);
27 |     void on_new_topic(const std::string& topic);
28 |     void on_consumer_commit(const std::string& group_id, const std::string& topic,
29 |                             int partition, int64_t offset);
30 |     void on_topic_message(const std::string& topic, int partition, int64_t offset);
31 |     void consumer_subscribe(const std::string& group_id);
32 |     void topic_subscribe(const std::string& topic);
33 | 
34 |     std::shared_ptr<OffsetStore> offset_store_;
35 |     bool track_consumer_commits_{false};
36 |     bool track_topic_messages_{false};
37 | };
38 | 
39 | } // api
40 | } // pirulo
41 | 


--------------------------------------------------------------------------------
/include/utils/async_observer.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | #include "utils/thread_pool.h"
 4 | #include "utils/observer.h"
 5 | 
 6 | namespace pirulo {
 7 | 
 8 | template <typename T, typename... Args>
 9 | class AsyncObserver {
10 | public:
11 |     using ObserverCallback = typename Observer<T, Args...>::ObserverCallback;
12 | 
13 |     AsyncObserver(ThreadPool& pool);
14 |     AsyncObserver(ThreadPool& pool, std::chrono::milliseconds cool_down_time);
15 | 
16 |     void observe(const T& object, const ObserverCallback& callback);
17 |     void notify(const T& object, const Args&... args);
18 | private:
19 |     Observer<T, Args...> observer_;
20 |     ThreadPool& pool_;
21 | };
22 | 
23 | template <typename T, typename... Args>
24 | AsyncObserver<T, Args...>::AsyncObserver(ThreadPool& pool)
25 | : pool_(pool) {
26 | 
27 | }
28 | 
29 | template <typename T, typename... Args>
30 | AsyncObserver<T, Args...>::AsyncObserver(ThreadPool& pool,
31 |                                          std::chrono::milliseconds cool_down_time)
32 | : observer_(cool_down_time), pool_(pool) {
33 | 
34 | }
35 | 
36 | template <typename T, typename... Args>
37 | void AsyncObserver<T, Args...>::observe(const T& object, const ObserverCallback& callback) {
38 |     observer_.observe(object, [&, callback](const T& object, const Args&... args) {
39 |         auto wrapped_callback = std::bind(callback, object, args...);
40 |         pool_.add_task([wrapped_callback]() {
41 |             wrapped_callback();
42 |         });
43 |     });
44 | }
45 | 
46 | template <typename T, typename... Args>
47 | void AsyncObserver<T, Args...>::notify(const T& object, const Args&... args) {
48 |     observer_.notify(object, args...);
49 | }
50 | 
51 | } // pirulo
52 | 


--------------------------------------------------------------------------------
/src/application.cpp:
--------------------------------------------------------------------------------
 1 | #include <thread>
 2 | #include <vector>
 3 | #include <cassert>
 4 | #include "application.h"
 5 | #include "detail/logging.h"
 6 | 
 7 | using std::thread;
 8 | using std::vector;
 9 | 
10 | namespace pirulo {
11 | 
12 | PIRULO_CREATE_LOGGER("p.app");
13 | 
14 | Application::Application(TopicOffsetReaderPtr topic_reader,
15 |                          ConsumerOffsetReaderPtr consumer_reader)
16 | : topic_reader_(move(topic_reader)), consumer_reader_(move(consumer_reader)) {
17 |     // Ensure everything uses the same store
18 |     assert(topic_reader_->get_store() == consumer_reader_->get_store());
19 | }
20 | 
21 | void Application::run() {
22 |     auto on_consumer_offset_eof = [&] {
23 |         const auto store = consumer_reader_->get_store();
24 |         // Enable notifications on the offset store
25 |         store->enable_notifications();
26 | 
27 |         LOG4CXX_INFO(logger, "Initializing " << plugins_.size() << " plugins");
28 |         // When we finish loading the __consumer_offsets topic, launch all plugins
29 |         for (auto& plugin_ptr : plugins_) {
30 |             plugin_ptr->launch(store);
31 |         }
32 |     };
33 | 
34 |     // Start topic and consumer offset consumption
35 |     vector<thread> threads;
36 |     threads.emplace_back([&] {
37 |         topic_reader_->run();
38 |     });
39 |     threads.emplace_back([&] {
40 |         consumer_reader_->run(on_consumer_offset_eof);
41 |     });
42 | 
43 |     // wut
44 |     process();
45 | 
46 |     for (thread& th : threads) {
47 |         th.join();
48 |     }
49 | }
50 | 
51 | void Application::stop() {
52 |     topic_reader_->stop();
53 |     consumer_reader_->stop();
54 | }
55 | 
56 | void Application::add_plugin(PluginPtr plugin) {
57 |     plugins_.emplace_back(move(plugin));
58 | }
59 | 
60 | void Application::process() {
61 | 
62 | }
63 | 
64 | } // pirulo
65 | 


--------------------------------------------------------------------------------
/src/python/helpers.cpp:
--------------------------------------------------------------------------------
 1 | #include <atomic>
 2 | #include <boost/python/object.hpp>
 3 | #include <boost/python/str.hpp>
 4 | #include <boost/python/import.hpp>
 5 | #include <boost/python/extract.hpp>
 6 | #include "python/helpers.h"
 7 | #include "python/api.h"
 8 | #include "offset_store.h"
 9 | #include "detail/logging.h"
10 | 
11 | using std::string;
12 | using std::vector;
13 | using std::call_once;
14 | using std::once_flag;
15 | 
16 | namespace python = boost::python;
17 | 
18 | namespace pirulo {
19 | namespace api {
20 | namespace helpers {
21 | 
22 | PIRULO_CREATE_LOGGER("p.python");
23 | 
24 | GILAcquirer::GILAcquirer()
25 | : gstate_(PyGILState_Ensure()) {
26 | }
27 | 
28 | GILAcquirer::~GILAcquirer() {
29 |     PyGILState_Release(gstate_);
30 | }
31 | 
32 | // Taken from https://stackoverflow.com/questions/1418015/how-to-get-python-exception-text
33 | string format_python_exception() {
34 |     using namespace boost::python;
35 |     using namespace boost;
36 | 
37 |     PyObject *exc, *val, *tb;
38 |     object formatted_list, formatted;
39 |     PyErr_Fetch(&exc, &val, &tb);
40 |     handle<> hexc(exc);
41 |     handle<> hval(allow_null(val));
42 |     handle<> htb(allow_null(tb)); 
43 |     object traceback(import("traceback"));
44 |     if (!tb) {
45 |         object format_exception_only(traceback.attr("format_exception_only"));
46 |         formatted_list = format_exception_only(hexc, hval);
47 |     } else {
48 |         object format_exception(traceback.attr("format_exception"));
49 |         formatted_list = format_exception(hexc, hval, htb);
50 |     }
51 |     formatted = str("\n").join(formatted_list);
52 |     return extract<string>(formatted);
53 | }
54 | 
55 | void initialize_python() {
56 |     static once_flag flag;
57 |     call_once(flag, [&] {
58 |         register_module();
59 |         PyEval_InitThreads();
60 |         Py_Initialize();
61 |         register_types();
62 |         PyEval_SaveThread();
63 |     });
64 | }
65 | 
66 | } // helpers
67 | } // api
68 | } // pirulo
69 | 


--------------------------------------------------------------------------------
/include/topic_offset_reader.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | #include <mutex>
 4 | #include <set>
 5 | #include <chrono>
 6 | #include <cppkafka/consumer.h>
 7 | #include "utils/thread_pool.h"
 8 | #include "utils/task_scheduler.h"
 9 | #include "offset_store.h"
10 | #include "consumer_offset_reader.h"
11 | #include "consumer_pool.h"
12 | 
13 | namespace pirulo {
14 | 
15 | class TopicOffsetReader {
16 | public:
17 |     using StorePtr = std::shared_ptr<OffsetStore>;
18 |     using ConsumerOffsetReaderPtr = std::shared_ptr<ConsumerOffsetReader>;
19 | 
20 |     TopicOffsetReader(StorePtr store, size_t thread_count,
21 |                       ConsumerOffsetReaderPtr consumer_reader,
22 |                       cppkafka::Configuration config);
23 | 
24 |     void run();
25 |     void stop();
26 | 
27 |     StorePtr get_store() const;
28 | private:
29 |     using TopicPartitionCount = std::unordered_map<std::string, size_t>;
30 |     using MetadataCallback = std::function<void(TopicPartitionCount)>;
31 |     using TopicTaskIdMap = std::map<cppkafka::TopicPartition, TaskScheduler::TaskId>;
32 | 
33 |     void async_process_topics(const TopicPartitionCount& topics);
34 |     void monitor_topics(const TopicPartitionCount& topics);
35 |     void monitor_topics(const cppkafka::TopicPartitionList& topics);
36 |     void monitor_new_topics();
37 |     cppkafka::TopicPartitionList get_new_topic_partitions(const TopicPartitionCount& counts);
38 |     TopicPartitionCount load_metadata();
39 |     void process_metadata(const MetadataCallback& callback);
40 |     void process_topic_partition(const cppkafka::TopicPartition& topic_partition);
41 |     void on_commit(const std::string& topic, int partition);
42 | 
43 |     ConsumerPool consumer_pool_;
44 |     StorePtr store_;
45 |     ThreadPool thread_pool_;
46 |     TaskScheduler task_scheduler_;
47 |     ConsumerOffsetReaderPtr consumer_offset_reader_;
48 |     TopicTaskIdMap monitored_topic_task_id_;
49 |     std::chrono::seconds maximum_topic_reload_time_{100};
50 |     std::chrono::seconds maximum_metadata_reload_time_{100};
51 |     bool running_{true};
52 | };
53 | 
54 | } // pirulo
55 | 


--------------------------------------------------------------------------------
/include/utils/task_scheduler.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | #include <functional>
 4 | #include <unordered_map>
 5 | #include <mutex>
 6 | #include <condition_variable>
 7 | #include <chrono>
 8 | #include <thread>
 9 | #include <atomic>
10 | #include <deque>
11 | 
12 | namespace pirulo {
13 | 
14 | class TaskScheduler {
15 | public:
16 |     using Task = std::function<void()>;
17 |     using TaskId = size_t;
18 |     using Duration = std::chrono::milliseconds;
19 | 
20 |     TaskScheduler();
21 |     TaskScheduler(const TaskScheduler&) = delete;
22 |     TaskScheduler& operator=(const TaskScheduler&) = delete;
23 |     ~TaskScheduler();
24 | 
25 |     TaskId add_task(Task task, Duration maximum_offset);
26 |     void remove_task(TaskId id);
27 |     void set_priority(TaskId id, double priority);
28 |     void set_minimum_reschedule_time(Duration value);
29 | 
30 | private:
31 |     using ClockType = std::chrono::steady_clock;
32 |     struct TaskExecutionInstance {
33 |         TaskId task_id;
34 |         ClockType::time_point scheduled_at;
35 |         ClockType::time_point scheduled_for;
36 |     };
37 |     using TaskQueue = std::deque<TaskExecutionInstance>;
38 |     struct TaskMetadata {
39 |         Task task;
40 |         TaskQueue::iterator queue_iterator;
41 |         ClockType::time_point last_priority_set_time;
42 |         Duration maximum_offset;
43 |         double priority;
44 |     };
45 |     using TaskMap = std::unordered_map<TaskId, TaskMetadata>;
46 | 
47 |     static Duration get_schedule_delta(const TaskMetadata& meta);
48 |     void stop();
49 |     TaskQueue::iterator schedule_task(TaskId task_id, const TaskMetadata& meta);
50 |     void remove_scheduled_instance(TaskId task_id);
51 |     void process_tasks();
52 | 
53 |     TaskMap tasks_;
54 |     TaskId current_task_id_{0};
55 |     std::thread process_thread_;
56 |     Duration minimum_reschedule_ = std::chrono::seconds(10);
57 |     Duration priority_adjustment_offset_ = std::chrono::seconds(60);
58 |     TaskQueue tasks_queue_;
59 |     mutable std::mutex tasks_mutex_;
60 |     std::condition_variable tasks_condition_;
61 |     std::atomic<bool> running_{true};
62 | };
63 | 
64 | } // pirulo
65 | 


--------------------------------------------------------------------------------
/src/utils/thread_pool.cpp:
--------------------------------------------------------------------------------
 1 | #include "utils/thread_pool.h"
 2 | 
 3 | using std::mutex;
 4 | using std::lock_guard;
 5 | using std::unique_lock;
 6 | using std::move;
 7 | 
 8 | namespace pirulo {
 9 | 
10 | ThreadPool::ThreadPool(size_t thread_count)
11 | : ThreadPool(thread_count, 0) {
12 | 
13 | }
14 | 
15 | ThreadPool::ThreadPool(size_t thread_count, size_t maximum_tasks)
16 | : maximum_tasks_(maximum_tasks) {
17 |     for (size_t i = 0; i < thread_count; ++i) {
18 |         threads_.emplace_back(&ThreadPool::process, this);
19 |     }
20 | }
21 | 
22 | ThreadPool::~ThreadPool() {
23 |     stop();
24 | }
25 | 
26 | bool ThreadPool::add_task(Task task) {
27 |     lock_guard<mutex> _(tasks_mutex_);
28 |     if (maximum_tasks_ > 0 && tasks_.size() >= maximum_tasks_) {
29 |         return false;
30 |     }
31 |     tasks_.push(move(task));
32 |     tasks_condition_.notify_one();
33 |     return true;
34 | }
35 | 
36 | void ThreadPool::stop() {
37 |     {
38 |         // Wake all threads
39 |         running_ = false;
40 |         lock_guard<mutex> _(tasks_mutex_);
41 |         tasks_condition_.notify_all();
42 |     }
43 |     for (auto& thread : threads_) {
44 |         thread.join();
45 |     }
46 |     threads_.clear();
47 | }
48 | 
49 | void ThreadPool::wait_for_tasks() {
50 |     while (running_) {
51 |         unique_lock<mutex> lock(tasks_mutex_);
52 |         if (tasks_.empty()) {
53 |             return;
54 |         }
55 |         no_tasks_condition_.wait(lock);
56 |     }
57 | }
58 | 
59 | void ThreadPool::process() {
60 |     while (running_) {
61 |         unique_lock<mutex> lock(tasks_mutex_);
62 | 
63 |         // Either spurious wake up or someone called ThreadPool::stop
64 |         while (running_ && tasks_.empty()) {
65 |             tasks_condition_.wait(lock);
66 |         }
67 |         if (!running_) {
68 |             continue;
69 |         }
70 | 
71 |         Task task = move(tasks_.front());
72 |         tasks_.pop();
73 | 
74 |         if (tasks_.empty()) {
75 |             no_tasks_condition_.notify_all();
76 |         }
77 | 
78 |         // Release lock and execute task
79 |         lock.unlock();
80 |         task();
81 |     }
82 | }
83 | 
84 | } // pirulo
85 | 


--------------------------------------------------------------------------------
/plugins/logger.py:
--------------------------------------------------------------------------------
 1 | import web
 2 | import threading
 3 | import time
 4 | import sys
 5 | from pirulo import Handler
 6 | 
 7 | class TopicsHandler:
 8 |     def GET(self):
 9 |         return Plugin.INSTANCE.topics
10 | 
11 | class ConsumersHandler:
12 |     def GET(self):
13 |         return Plugin.INSTANCE.consumers
14 | 
15 | class Plugin(Handler):
16 |     INSTANCE = None
17 | 
18 |     def __init__(self):
19 |         Handler.__init__(self)
20 |         Plugin.INSTANCE = self
21 |         self.fd = open('/tmp/events', 'w')
22 |         self.topics = []
23 |         self.consumers = []
24 |         self.thread = threading.Thread(target=self.launch_server)
25 |         self.thread.start()
26 | 
27 |     def handle_initialize(self):
28 |         self.subscribe_to_consumers()
29 |         self.subscribe_to_consumer_commits()
30 |         self.subscribe_to_topics()
31 |         self.subscribe_to_topic_message()
32 | 
33 |     def launch_server(self):
34 |         try:
35 |             urls = (
36 |                 '/topics', 'TopicsHandler',
37 |                 '/consumers', 'ConsumersHandler',
38 |             )
39 |             sys.argv = []
40 |             app = web.application(urls, globals())
41 |             app.run()
42 |         except Exception as ex:
43 |             print 'Failed running server: ' + str(ex)
44 | 
45 |     def log_message(self, message):
46 |         self.fd.write(message + '\n')
47 |         self.fd.flush()
48 | 
49 |     def handle_new_consumer(self, group_id):
50 |         self.log_message('New consumer {0} found'.format(group_id))
51 |         self.consumers.append(group_id)
52 | 
53 |     def handle_new_topic(self, topic):
54 |         self.log_message('Found topic {0}'.format(topic))
55 |         self.topics.append(topic)
56 | 
57 |     def handle_consumer_commit(self, group_id, topic, partition, offset):
58 |         self.log_message('Consumer {0} committed to {1}/{2} offset {3}'.format(
59 |             group_id,
60 |             topic,
61 |             partition,
62 |             offset 
63 |         ))
64 | 
65 |     def handle_topic_message(self, topic, partition, offset):
66 |         self.log_message('New offset for topic {0}/{1} at offset {2}'.format(
67 |             topic,
68 |             partition,
69 |             offset
70 |         ))
71 | 
72 | def create_plugin():
73 |     return Plugin()
74 | 


--------------------------------------------------------------------------------
/include/utils/observer.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | #include <map>
 4 | #include <chrono>
 5 | #include <functional>
 6 | #include <mutex>
 7 | 
 8 | namespace pirulo {
 9 | 
10 | template <typename T, typename... Args>
11 | class Observer {
12 | public:
13 |     using ObserverCallback = std::function<void(const T&, const Args&...)>;
14 | 
15 |     Observer();
16 |     Observer(std::chrono::milliseconds cool_down_time);
17 | 
18 |     void observe(const T& object, ObserverCallback callback);
19 |     void notify(const T& object, const Args&... args);
20 | 
21 | private:
22 |     using ClockType = std::chrono::steady_clock;
23 |     struct ObservedContext {
24 |         std::vector<ObserverCallback> observers;
25 |         ClockType::time_point last_observe_time;
26 |     };
27 |     using ObservedObjectsMap = std::map<T, ObservedContext>;
28 | 
29 |     ObservedObjectsMap observed_objects_;
30 |     std::chrono::milliseconds cool_down_time_;
31 |     mutable std::mutex observed_objects_mutex_;
32 | };
33 | 
34 | template <typename T, typename... Args>
35 | Observer<T, Args...>::Observer()
36 | : cool_down_time_(0) {
37 |     
38 | }
39 | 
40 | template <typename T, typename... Args>
41 | Observer<T, Args...>::Observer(std::chrono::milliseconds cool_down_time)
42 | : cool_down_time_(cool_down_time) {
43 | 
44 | }
45 | 
46 | template <typename T, typename... Args>
47 | void Observer<T, Args...>::observe(const T& object, ObserverCallback callback) {
48 |     std::lock_guard<std::mutex> _(observed_objects_mutex_);
49 |     observed_objects_[object].observers.emplace_back(std::move(callback));
50 | }
51 | 
52 | template <typename T, typename... Args>
53 | void Observer<T, Args...>::notify(const T& object, const Args&... args) {
54 |     std::unique_lock<std::mutex> lock(observed_objects_mutex_);
55 |     auto iter = observed_objects_.find(object);
56 |     if (iter == observed_objects_.end()) {
57 |         return;
58 |     }
59 |     auto now = ClockType::now();
60 |     // If we're still in cooldown phase, then don't trigger any callbacks
61 |     if (iter->second.last_observe_time + cool_down_time_ > now) {
62 |         return;
63 |     }
64 |     iter->second.last_observe_time = now;
65 | 
66 |     // Get the observers and release the lock
67 |     const std::vector<ObserverCallback> observers = iter->second.observers;
68 |     lock.unlock();
69 | 
70 |     for (const ObserverCallback& callback : observers) {
71 |         callback(object, args...);
72 |     }
73 | }
74 | 
75 | } // pirulo
76 | 


--------------------------------------------------------------------------------
/executables/topic_offsets.cpp:
--------------------------------------------------------------------------------
 1 | #include <iostream>
 2 | #include <thread>
 3 | #include <stdexcept>
 4 | #include <boost/program_options.hpp>
 5 | #include "topic_offset_reader.h"
 6 | #include "detail/logging.h"
 7 | 
 8 | using std::cout;
 9 | using std::endl;
10 | using std::move;
11 | using std::string;
12 | using std::make_shared;
13 | using std::thread;
14 | using std::cin;
15 | using std::exception;
16 | 
17 | using boost::optional;
18 | 
19 | using cppkafka::Configuration;
20 | using cppkafka::TopicPartition;
21 | 
22 | using pirulo::TopicOffsetReader;
23 | using pirulo::OffsetStore;
24 | using pirulo::logging::register_console_logger;
25 | 
26 | namespace po = boost::program_options;
27 | 
28 | int main(int argc, char* argv[]) {
29 |     string brokers;
30 |     string group_id;
31 | 
32 |     po::options_description options("Options");
33 |     options.add_options()
34 |         ("help,h",       "produce this help message")
35 |         ("brokers,b",    po::value<string>(&brokers)->required(), 
36 |                          "the kafka broker list")
37 |         ("group-id,g",   po::value<string>(&group_id),
38 |                          "the consumer group id to be used")
39 |         ;
40 | 
41 |     po::variables_map vm;
42 | 
43 |     try {
44 |         po::store(po::command_line_parser(argc, argv).options(options).run(), vm);
45 |         po::notify(vm);
46 |     }
47 |     catch (const exception& ex) {
48 |         cout << "Error parsing options: " << ex.what() << endl;
49 |         cout << endl;
50 |         cout << options << endl;
51 |         return 1;
52 |     }
53 | 
54 |     register_console_logger();
55 | 
56 |     // Construct the configuration
57 |     Configuration config = {
58 |         { "metadata.broker.list", brokers },
59 |         { "group.id", group_id },
60 |         // Disable auto commit
61 |         { "enable.auto.commit", false },
62 |     };
63 | 
64 |     auto store = make_shared<OffsetStore>();
65 |     TopicOffsetReader reader(store, 2, move(config));
66 |     thread th([&]() {
67 |         reader.run();
68 |     });
69 | 
70 |     string topic;
71 |     int partition;
72 |     while (cin >> topic >> partition) {
73 |         optional<int64_t> offset = store->get_topic_offset(topic, partition);
74 |         if (!offset) {
75 |             cout << "Topic/partition not found\n";
76 |         }
77 |         else {
78 |             cout << "Offset for " << TopicPartition(topic, partition) << ": " << *offset << endl;
79 |         }
80 |     }
81 | 
82 |     reader.stop();
83 |     th.join();
84 | }


--------------------------------------------------------------------------------
/executables/consumer_offsets.cpp:
--------------------------------------------------------------------------------
 1 | #include <iostream>
 2 | #include <thread>
 3 | #include <stdexcept>
 4 | #include <boost/program_options.hpp>
 5 | #include "consumer_offset_reader.h"
 6 | #include "detail/logging.h"
 7 | 
 8 | using std::cout;
 9 | using std::endl;
10 | using std::move;
11 | using std::string;
12 | using std::make_shared;
13 | using std::thread;
14 | using std::cin;
15 | using std::exception;
16 | 
17 | using std::chrono::seconds;
18 | 
19 | using cppkafka::Configuration;
20 | 
21 | using pirulo::ConsumerOffsetReader;
22 | using pirulo::OffsetStore;
23 | using pirulo::logging::register_console_logger;
24 | 
25 | namespace po = boost::program_options;
26 | 
27 | int main(int argc, char* argv[]) {
28 |     string brokers;
29 |     string group_id;
30 | 
31 |     po::options_description options("Options");
32 |     options.add_options()
33 |         ("help,h",       "produce this help message")
34 |         ("brokers,b",    po::value<string>(&brokers)->required(), 
35 |                          "the kafka broker list")
36 |         ("group-id,g",   po::value<string>(&group_id),
37 |                          "the consumer group id to be used")
38 |         ;
39 | 
40 |     po::variables_map vm;
41 | 
42 |     try {
43 |         po::store(po::command_line_parser(argc, argv).options(options).run(), vm);
44 |         po::notify(vm);
45 |     }
46 |     catch (const exception& ex) {
47 |         cout << "Error parsing options: " << ex.what() << endl;
48 |         cout << endl;
49 |         cout << options << endl;
50 |         return 1;
51 |     }
52 | 
53 |     register_console_logger();
54 | 
55 |     // Construct the configuration
56 |     Configuration config = {
57 |         { "metadata.broker.list", brokers },
58 |         { "group.id", group_id },
59 |         // Disable auto commit
60 |         { "enable.auto.commit", false }
61 |     };
62 | 
63 |     auto store = make_shared<OffsetStore>();
64 |     auto on_eof = [&]() {
65 |         cout << "Reached EOF on all partitions\n";
66 |     };
67 |     ConsumerOffsetReader reader(store, seconds(10), move(config));
68 |     thread th([&]() {
69 |         reader.run(on_eof);
70 |     });
71 | 
72 |     string consumer_group;
73 |     while (cin >> consumer_group) {
74 |         auto offsets = store->get_consumer_offsets(consumer_group);
75 |         if (offsets.empty()) {
76 |             cout << "Consumer not found\n";
77 |         }
78 |         else {
79 |             for (const auto& offset : offsets) {
80 |                 cout << offset.get_topic_partition() << ": "
81 |                      << offset.get_topic_partition().get_offset() << endl;
82 |             }
83 |         }
84 |     }
85 | 
86 |     reader.stop();
87 |     th.join();
88 | }


--------------------------------------------------------------------------------
/src/python/plugin.cpp:
--------------------------------------------------------------------------------
 1 | #include <stdexcept>
 2 | #include <boost/python/object.hpp>
 3 | #include <boost/python/import.hpp>
 4 | #include <boost/python/exec.hpp>
 5 | #include <boost/python/dict.hpp>
 6 | #include "python/plugin.h"
 7 | #include "python/helpers.h"
 8 | #include "detail/logging.h"
 9 | #include "offset_store.h"
10 | 
11 | using std::string;
12 | using std::vector;
13 | using std::once_flag;
14 | using std::call_once;
15 | using std::runtime_error;
16 | 
17 | using boost::optional;
18 | 
19 | namespace python = boost::python;
20 | 
21 | namespace pirulo {
22 | namespace api {
23 | 
24 | PIRULO_CREATE_LOGGER("p.python");
25 | 
26 | PythonPlugin::PythonPlugin(const string& modules_path, const string& file_path) {
27 |     helpers::initialize_python();
28 |     helpers::GILAcquirer _;
29 | 
30 |     // Dirtiness taken from https://wiki.python.org/moin/boost.python/EmbeddingPython
31 |     python::dict locals;
32 | 
33 |     python::object main_module = python::import("__main__");
34 |     python::object main_namespace = main_module.attr("__dict__");
35 | 
36 |     locals["modules_path"] = modules_path;
37 |     locals["path"] = file_path;
38 |     try {
39 |         python::exec("import imp, sys, os.path\n"
40 |              "sys.path.append(os.path.abspath(modules_path))\n"
41 |              "module_name = os.path.basename(path)[:-3]\n"
42 |              "module = imp.load_module(module_name,open(path),path,('py','U',imp.PY_SOURCE))\n",
43 |              main_namespace, locals);
44 |     }
45 |     catch (const python::error_already_set& ex) {
46 |         LOG4CXX_ERROR(logger, "Failed to load module " << file_path
47 |                       << ": " << helpers::format_python_exception());
48 |         throw runtime_error("Error loading python plugin " + file_path);
49 |     }
50 | 
51 |     try {
52 |         python::object plugin_factory = locals["module"].attr("create_plugin");
53 |         plugin_ = plugin_factory();
54 |     }
55 |     catch (const python::error_already_set& ex) {
56 |         LOG4CXX_ERROR(logger, "Error instantiating Plugin classs on module " << file_path
57 |                       << ": " << helpers::format_python_exception());
58 |         throw runtime_error("Error instanting python plugin " + file_path);
59 |     }
60 | }
61 | 
62 | PythonPlugin::~PythonPlugin() {
63 |     helpers::GILAcquirer _;
64 |     plugin_ = {};
65 | }
66 | 
67 | void PythonPlugin::initialize() {
68 |     helpers::GILAcquirer _;
69 |     try {
70 |         const StorePtr& store = get_store();
71 |         plugin_.attr("initialize")(store);
72 |     }
73 |     catch (const python::error_already_set& ex) {
74 |         LOG4CXX_ERROR(logger, "Error initializing plugin: "
75 |                       << helpers::format_python_exception());
76 |     }
77 | }
78 | 
79 | } // api
80 | } // pirulo
81 | 


--------------------------------------------------------------------------------
/include/detail/memory.h:
--------------------------------------------------------------------------------
  1 | // Code taken from libtins: https://github.com/mfontanini/libtins/blob/master/include/tins/memory_helpers.h
  2 | 
  3 | #pragma once
  4 | 
  5 | #include <cstdint>
  6 | #include <cstring>
  7 | #include <string>
  8 | #include <cppkafka/buffer.h>
  9 | #include "../exceptions.h"
 10 | #include "endianness.h"
 11 | 
 12 | namespace pirulo {
 13 | 
 14 | inline void read_data(const uint8_t* buffer, uint8_t* output_buffer, size_t size) {
 15 |     std::memcpy(output_buffer, buffer, size);
 16 | }
 17 | 
 18 | template <typename T>
 19 | void read_value(const uint8_t* buffer, T& value) {
 20 |     std::memcpy(&value, buffer, sizeof(value));
 21 | }
 22 | 
 23 | class InputMemoryStream {
 24 | public:
 25 |     InputMemoryStream(const uint8_t* buffer, size_t total_sz)
 26 |     : buffer_(buffer), size_(total_sz) {
 27 |     }
 28 | 
 29 |     InputMemoryStream(const cppkafka::Buffer& data)
 30 |     : buffer_(data.get_data()), size_(data.get_size()) {
 31 |     }
 32 |  
 33 |     template <typename T>
 34 |     T read() {
 35 |         T output;
 36 |         read(output);
 37 |         return output;
 38 |     }
 39 | 
 40 |     template <typename T>
 41 |     T read_le() {
 42 |         return endian::le_to_host(read<T>());
 43 |     }
 44 | 
 45 |     template <typename T>
 46 |     T read_be() {
 47 |         return endian::be_to_host(read<T>());
 48 |     }
 49 | 
 50 |     template <typename T>
 51 |     void read(T& value) {
 52 |         if (!can_read(sizeof(value))) {
 53 |             throw ParseException();
 54 |         }
 55 |         read_value(buffer_, value);
 56 |         skip(sizeof(value));
 57 |     }
 58 | 
 59 |     void read(std::string& value) {
 60 |         uint16_t length = read_be<uint16_t>();
 61 |         if (!can_read(length)) {
 62 |             throw ParseException();
 63 |         }
 64 |         value.assign(pointer(), pointer() + length);
 65 |         skip(length);
 66 |     }
 67 | 
 68 |     void skip(size_t size) {
 69 |         if (size > size_) {
 70 |             throw ParseException();
 71 |         }
 72 |         buffer_ += size;
 73 |         size_ -= size;
 74 |     }
 75 | 
 76 |     bool can_read(size_t byte_count) const {
 77 |         return size_ >= byte_count;
 78 |     }
 79 | 
 80 |     void read(void* output_buffer, size_t output_buffer_size) {
 81 |         if (!can_read(output_buffer_size)) {
 82 |             throw ParseException();
 83 |         }
 84 |         read_data(buffer_, (uint8_t*)output_buffer, output_buffer_size);
 85 |         skip(output_buffer_size);
 86 |     }
 87 | 
 88 |     const uint8_t* pointer() const {
 89 |         return buffer_;
 90 |     }
 91 | 
 92 |     size_t size() const {
 93 |         return size_;
 94 |     }
 95 | 
 96 |     void size(size_t new_size) {
 97 |         size_ = new_size;
 98 |     }
 99 | 
100 |     operator bool() const {
101 |         return size_ > 0;
102 |     }
103 | private:
104 |     const uint8_t* buffer_;
105 |     size_t size_;
106 | };
107 | 
108 | } // pirulo
109 | 


--------------------------------------------------------------------------------
/src/main.cpp:
--------------------------------------------------------------------------------
 1 | #include <iostream>
 2 | #include <thread>
 3 | #include <stdexcept>
 4 | #include <csignal>
 5 | #include <boost/program_options.hpp>
 6 | #include "application.h"
 7 | #include "python/plugin.h"
 8 | #include "detail/logging.h"
 9 | 
10 | using std::cout;
11 | using std::endl;
12 | using std::move;
13 | using std::string;
14 | using std::make_shared;
15 | using std::thread;
16 | using std::cin;
17 | using std::exception;
18 | using std::unique_ptr;
19 | using std::function;
20 | 
21 | using std::chrono::seconds;
22 | 
23 | using cppkafka::Configuration;
24 | 
25 | using pirulo::Application;
26 | using pirulo::ConsumerOffsetReader;
27 | using pirulo::TopicOffsetReader;
28 | using pirulo::OffsetStore;
29 | using pirulo::api::PythonPlugin;
30 | using pirulo::logging::register_console_logger;
31 | 
32 | namespace po = boost::program_options;
33 | 
34 | function<void()> signal_handler;
35 | 
36 | void on_signal(int) {
37 |     signal_handler();
38 | }
39 | 
40 | int main(int argc, char* argv[]) {
41 |     string brokers;
42 |     string group_id;
43 |     unsigned threads;
44 | 
45 |     po::options_description options("Options");
46 |     options.add_options()
47 |         ("help,h",       "produce this help message")
48 |         ("brokers,b",    po::value<string>(&brokers)->required(), 
49 |                          "the kafka broker list")
50 |         ("threads,t",    po::value<unsigned>(&threads)->default_value(2),
51 |                          "amount of threads to use for topic metadata reloading")
52 |         ;
53 | 
54 |     po::variables_map vm;
55 | 
56 |     try {
57 |         po::store(po::command_line_parser(argc, argv).options(options).run(), vm);
58 |         po::notify(vm);
59 |     }
60 |     catch (const exception& ex) {
61 |         cout << "Error parsing options: " << ex.what() << endl;
62 |         cout << endl;
63 |         cout << options << endl;
64 |         return 1;
65 |     }
66 | 
67 |     register_console_logger("", "TRACE");
68 | 
69 |     // Construct the configuration
70 |     Configuration config = {
71 |         { "metadata.broker.list", brokers },
72 |         // Disable auto commit
73 |         { "enable.auto.commit", false }
74 |     };
75 | 
76 |     auto store = make_shared<OffsetStore>();
77 |     auto consumer_reader = make_shared<ConsumerOffsetReader>(store, seconds(10), config);
78 |     auto topic_reader = make_shared<TopicOffsetReader>(store, threads, consumer_reader,
79 |                                                        config);
80 | 
81 |     Application app(move(topic_reader), move(consumer_reader));
82 |     // app.add_plugin(unique_ptr<PythonPlugin>(new PythonPlugin("../plugins",
83 |     //                                                         "../plugins/logger.py")));
84 |     app.add_plugin(unique_ptr<PythonPlugin>(new PythonPlugin("../plugins",
85 |                                                              "../plugins/lag_exporter.py")));
86 | 
87 |     signal_handler = [&] {
88 |         app.stop();
89 |     };
90 | 
91 |     signal(SIGINT, &on_signal);
92 |     signal(SIGTERM, &on_signal);
93 |     signal(SIGQUIT, &on_signal);
94 | 
95 |     app.run();
96 | }
97 | 


--------------------------------------------------------------------------------
/include/offset_store.h:
--------------------------------------------------------------------------------
 1 | #pragma once
 2 | 
 3 | #include <string>
 4 | #include <cstdint>
 5 | #include <unordered_map>
 6 | #include <unordered_set>
 7 | #include <map>
 8 | #include <mutex>
 9 | #include <vector>
10 | #include <functional>
11 | #include <boost/optional.hpp>
12 | #include <cppkafka/topic_partition.h>
13 | #include "consumer_offset.h"
14 | #include "utils/async_observer.h"
15 | #include "utils/thread_pool.h"
16 | 
17 | namespace pirulo {
18 | 
19 | class OffsetStore {
20 | public:
21 |     using ConsumerCallback = std::function<void(const std::string& group_id)>;
22 |     using TopicCallback = std::function<void(const std::string& topic_name)>;
23 |     using ConsumerCommitCallback = std::function<void(const std::string& group_id,
24 |                                                       const std::string& topic,
25 |                                                       int partition,
26 |                                                       uint64_t offset)>;
27 |     using TopicMessageCallback = std::function<void(const std::string& topic,
28 |                                                     int partition,
29 |                                                     uint64_t offset)>;
30 | 
31 |     OffsetStore();
32 | 
33 |     void store_consumer_offset(const std::string& group_id, const std::string& topic,
34 |                                int partition, uint64_t offset);
35 |     void store_topic_offset(const std::string& topic, int partition, uint64_t offset);
36 |     void on_new_consumer(ConsumerCallback callback);
37 |     void on_new_topic(TopicCallback callback);
38 |     void on_consumer_commit(const std::string& group_id, ConsumerCommitCallback callback);
39 |     void on_topic_message(const std::string& topic, TopicMessageCallback callback);
40 | 
41 |     void enable_notifications();
42 | 
43 |     std::vector<std::string> get_consumers() const;
44 |     std::vector<ConsumerOffset> get_consumer_offsets(const std::string& group_id) const;
45 |     boost::optional<int64_t> get_topic_offset(const std::string& topic,
46 |                                               int partition) const;
47 |     std::vector<std::string> get_topics() const;
48 | private:
49 |     // Make sure tasks won't start piling up
50 |     static constexpr size_t MAXIMUM_OBSERVER_TASKS = 10000;
51 | 
52 |     using TopicMap = std::map<cppkafka::TopicPartition, int64_t>;
53 |     using ConsumerMap = std::unordered_map<std::string, TopicMap>;
54 |     using StringSet = std::unordered_set<std::string>;
55 | 
56 |     ConsumerMap consumer_offsets_;
57 |     TopicMap topic_offsets_;
58 |     StringSet consumers_;
59 |     StringSet topics_;
60 |     ThreadPool thread_pool_{1, MAXIMUM_OBSERVER_TASKS};
61 |     AsyncObserver<int, std::string> new_string_observer_;
62 |     AsyncObserver<std::string, std::string, int, uint64_t> consumer_commit_observer_;
63 |     AsyncObserver<std::string, int, uint64_t> topic_message_observer_; 
64 |     std::string new_consumer_id_;
65 |     mutable std::mutex consumer_offsets_mutex_;
66 |     mutable std::mutex topic_offsets_mutex_;
67 |     bool notifications_enabled_{false};
68 | };
69 | 
70 | } // pirulo
71 | 


--------------------------------------------------------------------------------
/src/python/lag_tracker_handler.cpp:
--------------------------------------------------------------------------------
 1 | #include <algorithm>
 2 | #include "detail/logging.h"
 3 | #include "python/lag_tracker_handler.h"
 4 | 
 5 | using std::string;
 6 | using std::vector;
 7 | using std::max;
 8 | using std::make_tuple;
 9 | using std::shared_ptr;
10 | 
11 | using boost::optional;
12 | 
13 | namespace pirulo {
14 | namespace api {
15 | 
16 | PIRULO_CREATE_LOGGER("p.lag_tracker");
17 | 
18 | void LagTrackerHandler::handle_initialize() {
19 |     LOG4CXX_INFO(logger, "Initializing lag tracker handler");
20 |     const auto& offset_store = get_offset_store();
21 |     for (const string& consumer : offset_store->get_consumers()) {
22 |         const vector<ConsumerOffset> offsets = offset_store->get_consumer_offsets(consumer);
23 |         for (const ConsumerOffset& offset : offsets) {
24 |             const auto& topic_partition = offset.get_topic_partition();
25 |             const string& topic = topic_partition.get_topic();
26 |             const int partition = topic_partition.get_partition();
27 |             TopicPartitionInfo& info = topic_partition_info_[make_tuple(topic, partition)];
28 |             info.consumer_offsets.emplace(consumer, topic_partition.get_offset());
29 |             if (info.offset == -1) {
30 |                 const optional<int64_t> maybe_offset = offset_store->get_topic_offset(topic,
31 |                                                                                       partition);
32 |                 if (maybe_offset) {
33 |                     info.offset = *maybe_offset;
34 |                 }
35 |             }
36 |         }
37 |     }
38 | 
39 |     subscribe_to_topics();
40 |     subscribe_to_topic_message();
41 |     subscribe_to_consumers();
42 |     subscribe_to_consumer_commits();
43 | }
44 | 
45 | void LagTrackerHandler::handle_new_consumer(const string& group_id) {
46 |     Handler::handle_new_consumer(group_id);
47 | }
48 | 
49 | void LagTrackerHandler::handle_new_topic(const string& topic) {
50 |     Handler::handle_new_topic(topic);
51 | }
52 | 
53 | void LagTrackerHandler::handle_consumer_commit(const string& group_id, const string& topic,
54 |                                                int partition, int64_t offset) {
55 |     auto& info = topic_partition_info_[make_tuple(topic, partition)];
56 |     info.consumer_offsets[group_id] = offset;
57 |     if (info.offset != -1) {
58 |         handle_lag_update(topic, partition, group_id, max<int64_t>(0, info.offset - offset));
59 |     }
60 |     Handler::handle_consumer_commit(group_id, topic, partition, offset);
61 | }
62 | 
63 | void LagTrackerHandler::handle_topic_message(const string& topic, int partition, int64_t offset) {
64 |     auto& info = topic_partition_info_[make_tuple(topic, partition)];
65 |     info.offset = offset;
66 |     for (const auto& consumer_offset_pair : info.consumer_offsets) {
67 |         const string& group_id = consumer_offset_pair.first;
68 |         const uint64_t lag = max<int64_t>(0, offset - consumer_offset_pair.second);
69 |         handle_lag_update(topic, partition, group_id, lag);
70 |     }
71 |     Handler::handle_topic_message(topic, partition, offset);
72 | }
73 | 
74 | } // api
75 | } // pirulo
76 | 


--------------------------------------------------------------------------------
/src/python/handler.cpp:
--------------------------------------------------------------------------------
  1 | #include <functional>
  2 | #include "python/handler.h"
  3 | #include "detail/logging.h"
  4 | 
  5 | using std::string;
  6 | using std::shared_ptr;
  7 | using std::bind;
  8 | using std::vector;
  9 | 
 10 | using namespace std::placeholders;
 11 | 
 12 | namespace pirulo {
 13 | namespace api {
 14 | 
 15 | PIRULO_CREATE_LOGGER("p.handler");
 16 | 
 17 | void Handler::initialize(const shared_ptr<OffsetStore>& store) {
 18 |     offset_store_ = store;
 19 |     handle_initialize();
 20 | }
 21 | 
 22 | void Handler::subscribe_to_consumers() {
 23 |     offset_store_->on_new_consumer(bind(&Handler::on_new_consumer, this, _1));
 24 |     for (const string& group_id : offset_store_->get_consumers()) {
 25 |         on_new_consumer(group_id);
 26 |     }
 27 | }
 28 | 
 29 | void Handler::subscribe_to_consumer_commits() {
 30 |     if (track_consumer_commits_) {
 31 |         return;
 32 |     }
 33 |     track_consumer_commits_ = true;
 34 |     for (const string& group_id : offset_store_->get_consumers()) {
 35 |         consumer_subscribe(group_id);
 36 |     }
 37 | }
 38 | 
 39 | void Handler::subscribe_to_topics() {
 40 |     offset_store_->on_new_topic(bind(&Handler::on_new_topic, this, _1));
 41 |     for (const string& topic : offset_store_->get_topics()) {
 42 |         on_new_topic(topic);
 43 |     }
 44 | }
 45 | 
 46 | void Handler::subscribe_to_topic_message() {
 47 |     if (track_topic_messages_) {
 48 |         return;
 49 |     }
 50 |     track_topic_messages_ = true;
 51 |     for (const string& topic : offset_store_->get_topics()) {
 52 |         topic_subscribe(topic);
 53 |     }
 54 | }
 55 | 
 56 | const shared_ptr<OffsetStore>& Handler::get_offset_store() const {
 57 |     return offset_store_;
 58 | }
 59 | 
 60 | void Handler::handle_initialize() {
 61 | 
 62 | }
 63 | 
 64 | void Handler::handle_new_consumer(const string& group_id) {
 65 | 
 66 | }
 67 | 
 68 | void Handler::handle_new_topic(const string& topic) {
 69 | 
 70 | }
 71 | 
 72 | void Handler::handle_consumer_commit(const string& group_id, const string& topic,
 73 |                                      int partition, int64_t offset) {
 74 | 
 75 | }
 76 | 
 77 | void Handler::handle_topic_message(const string& topic, int partition, int64_t offset) {
 78 | 
 79 | }
 80 | 
 81 | void Handler::on_new_consumer(const string& group_id) {
 82 |     LOG4CXX_DEBUG(logger, "Found new consumer: " << group_id);
 83 |     handle_new_consumer(group_id);
 84 |     if (track_consumer_commits_) {
 85 |         LOG4CXX_DEBUG(logger, "Subscribing to consumer");
 86 |         consumer_subscribe(group_id);
 87 |         const vector<ConsumerOffset> offsets = get_offset_store()->get_consumer_offsets(group_id);
 88 |         for (const ConsumerOffset& offset : offsets) {
 89 |             const auto& topic_partition = offset.get_topic_partition();
 90 |             on_consumer_commit(group_id, topic_partition.get_topic(),
 91 |                                topic_partition.get_partition(), topic_partition.get_offset());
 92 |         }
 93 |     }
 94 | }
 95 | 
 96 | void Handler::on_new_topic(const string& topic) {
 97 |     LOG4CXX_DEBUG(logger, "Found new topic: " << topic);
 98 |     handle_new_topic(topic);
 99 |     if (track_topic_messages_) {
100 |         topic_subscribe(topic);
101 |     }
102 | }
103 | 
104 | void Handler::on_consumer_commit(const string& group_id, const string& topic,
105 |                                  int partition, int64_t offset) {
106 |     handle_consumer_commit(group_id, topic, partition, offset);
107 | }
108 | 
109 | void Handler::on_topic_message(const string& topic, int partition, int64_t offset) {
110 |     handle_topic_message(topic, partition, offset);
111 | }
112 | 
113 | void Handler::consumer_subscribe(const string& group_id) {
114 |     offset_store_->on_consumer_commit(group_id,
115 |                                       bind(&Handler::on_consumer_commit, this, _1, _2,
116 |                                            _3, _4));
117 | }
118 | 
119 | void Handler::topic_subscribe(const string& topic) {
120 |     offset_store_->on_topic_message(topic,
121 |                                     bind(&Handler::on_topic_message, this, _1, _2, _3));
122 | }
123 | 
124 | } // api
125 | } // pirulo
126 | 


--------------------------------------------------------------------------------
/src/consumer_offset_reader.cpp:
--------------------------------------------------------------------------------
  1 | #include <cstdint>
  2 | #include "consumer_offset_reader.h"
  3 | #include "exceptions.h"
  4 | #include "detail/memory.h"
  5 | #include "detail/logging.h"
  6 | #include "utils/utils.h"
  7 | 
  8 | using std::move;
  9 | using std::string;
 10 | 
 11 | using std::chrono::milliseconds;
 12 | 
 13 | using cppkafka::Configuration;
 14 | using cppkafka::ConsumerDispatcher;
 15 | using cppkafka::Message;
 16 | using cppkafka::TopicPartition;
 17 | using cppkafka::TopicPartitionList;
 18 | 
 19 | namespace pirulo {
 20 | 
 21 | PIRULO_CREATE_LOGGER("p.offsets");
 22 | 
 23 | static Configuration prepare_config(Configuration config) {
 24 |     config.set_default_topic_configuration({{ "auto.offset.reset", "smallest" }});
 25 |     config.set("group.id", utils::generate_group_id());
 26 |     LOG4CXX_INFO(logger, "Using consumer " << config.get("group.id") << " for "
 27 |                  << " consumer offset consumption");
 28 |     return config;
 29 | }
 30 | 
 31 | ConsumerOffsetReader::ConsumerOffsetReader(StorePtr store, milliseconds consumer_offset_cool_down,
 32 |                                            Configuration config) :
 33 |     store_(move(store)), consumer_(prepare_config(move(config))),
 34 |     observer_(consumer_offset_cool_down) {
 35 | 
 36 | }
 37 | 
 38 | void ConsumerOffsetReader::run(const EofCallback& callback) {
 39 |     consumer_.set_assignment_callback([&](const TopicPartitionList& topic_partitions) {
 40 |         for (const TopicPartition& topic_partition : topic_partitions) {
 41 |             pending_partitions_.emplace(topic_partition.get_partition());
 42 |         }
 43 |     });
 44 |     LOG4CXX_INFO(logger, "Starting loading consumer offsets");
 45 |     
 46 |     consumer_.subscribe({ "__consumer_offsets" });
 47 |     dispatcher_.run(
 48 |         [&](Message msg) {
 49 |             try {
 50 |                 if (msg.get_payload()) {
 51 |                     handle_message(msg);
 52 |                 }
 53 |             }
 54 |             catch (const ParseException&) {
 55 |                 LOG4CXX_WARN(logger, "Failed to parse consumer offset record");
 56 |             }
 57 |         },
 58 |         [&](ConsumerDispatcher::EndOfFile, const TopicPartition& topic_partition) {
 59 |             // If we reached EOF on all partitions, execute the EOF callback
 60 |             if (pending_partitions_.erase(topic_partition.get_partition()) &&
 61 |                 pending_partitions_.empty()) {
 62 |                 LOG4CXX_INFO(logger, "Finished loading consumer offsets");
 63 |                 callback();
 64 | 
 65 |                 // Enable notifications for new commits
 66 |                 notifications_enabled_ = true;
 67 |             }
 68 |         }
 69 |     );
 70 | }
 71 | 
 72 | void ConsumerOffsetReader::stop() {
 73 |     dispatcher_.stop();
 74 | }
 75 | 
 76 | void ConsumerOffsetReader::watch_commits(const string& topic, int partition,
 77 |                                          TopicCommitCallback callback) {
 78 |     auto wrapped_callback = [callback](const TopicPartition& topic_partition) {
 79 |         callback(topic_partition.get_topic(), topic_partition.get_partition());
 80 |     };
 81 |     observer_.observe({ topic, partition }, move(wrapped_callback));
 82 | }
 83 | 
 84 | ConsumerOffsetReader::StorePtr ConsumerOffsetReader::get_store() const {
 85 |     return store_;
 86 | }
 87 | 
 88 | void ConsumerOffsetReader::handle_message(const Message& msg) {
 89 |     InputMemoryStream key_input(msg.get_key());
 90 |     uint16_t version = key_input.read_be<uint16_t>();
 91 |     if (version > 1) {
 92 |         return;
 93 |     }
 94 |     string group_id = key_input.read<string>();
 95 |     string topic = key_input.read<string>();
 96 |     int partition = key_input.read_be<uint32_t>();
 97 | 
 98 |     InputMemoryStream value_input(msg.get_payload());
 99 |     // Value version
100 |     version = value_input.read_be<uint16_t>();
101 |     if (version > 1) {
102 |         throw ParseException();
103 |     }
104 |     uint64_t offset = value_input.read_be<uint64_t>();
105 |     store_->store_consumer_offset(group_id, topic, partition, offset);
106 | 
107 |     if (notifications_enabled_) {
108 |         observer_.notify({ topic, partition });
109 |     }
110 | }
111 | 
112 | } // pirulo
113 | 


--------------------------------------------------------------------------------
/src/offset_store.cpp:
--------------------------------------------------------------------------------
  1 | #include "offset_store.h"
  2 | 
  3 | using std::string;
  4 | using std::mutex;
  5 | using std::lock_guard;
  6 | using std::vector;
  7 | using std::move;
  8 | 
  9 | using std::chrono::seconds;
 10 | using std::chrono::milliseconds;
 11 | 
 12 | using boost::optional;
 13 | 
 14 | using cppkafka::TopicPartition;
 15 | 
 16 | namespace pirulo {
 17 | 
 18 | static const int NEW_CONSUMER_ID = 0;
 19 | static const int NEW_TOPIC_ID = 1;
 20 | 
 21 | // TODO: don't hardcode these constants
 22 | OffsetStore::OffsetStore()
 23 | : new_string_observer_(thread_pool_), consumer_commit_observer_(thread_pool_, seconds(10)),
 24 |   topic_message_observer_(thread_pool_, seconds(10)) {
 25 | 
 26 | }
 27 | 
 28 | void OffsetStore::store_consumer_offset(const string& group_id, const string& topic,
 29 |                                         int partition, uint64_t offset) {
 30 |     bool is_new_consumer = false;
 31 |     {
 32 |         lock_guard<mutex> _(consumer_offsets_mutex_);
 33 |         consumer_offsets_[group_id][{ topic, partition }] = offset;
 34 |         is_new_consumer = consumers_.insert(group_id).second;
 35 |     }
 36 |     // If notifications aren't enabled, we're done
 37 |     if (!notifications_enabled_) {
 38 |         return;
 39 |     }
 40 | 
 41 |     // Notify that there was a new commit for this consumer group
 42 |     consumer_commit_observer_.notify(group_id, topic, partition, offset);
 43 | 
 44 |     // If this is a new consumer group, notify
 45 |     if (is_new_consumer) {
 46 |         new_string_observer_.notify(NEW_CONSUMER_ID, group_id);
 47 |     }
 48 | }
 49 | 
 50 | void OffsetStore::store_topic_offset(const string& topic, int partition,
 51 |                                      uint64_t offset) {
 52 |     bool is_new_topic = false;
 53 |     bool is_new_offset = false;
 54 |     {
 55 |         lock_guard<mutex> _(topic_offsets_mutex_);
 56 |         int64_t& existing_offset = topic_offsets_[{topic, partition}];
 57 |         is_new_offset = existing_offset != static_cast<int64_t>(offset);
 58 |         is_new_topic = topics_.emplace(topic).second;
 59 |         existing_offset = offset;
 60 |     }
 61 |     // If notifications aren't enabled, we're done
 62 |     if (!notifications_enabled_) {
 63 |         return;
 64 |     }
 65 | 
 66 |     if (is_new_offset) {
 67 |         topic_message_observer_.notify(topic, partition, offset);
 68 |     }
 69 |     if (is_new_topic) {
 70 |         new_string_observer_.notify(NEW_TOPIC_ID, topic);
 71 |     }
 72 | }
 73 | 
 74 | void OffsetStore::on_new_consumer(ConsumerCallback callback) {
 75 |     new_string_observer_.observe(NEW_CONSUMER_ID, [=](int, const string& group_id) {
 76 |         callback(group_id);
 77 |     });
 78 | }
 79 | 
 80 | void OffsetStore::on_new_topic(TopicCallback callback) {
 81 |     new_string_observer_.observe(NEW_TOPIC_ID, [=](int, const string& topic) {
 82 |         callback(topic);
 83 |     });
 84 | }
 85 | 
 86 | void OffsetStore::on_consumer_commit(const string& group_id, ConsumerCommitCallback callback) {
 87 |     consumer_commit_observer_.observe(group_id, move(callback));
 88 | }
 89 | 
 90 | void OffsetStore::on_topic_message(const string& topic, TopicMessageCallback callback) {
 91 |     topic_message_observer_.observe(topic, move(callback));
 92 | }
 93 | 
 94 | void OffsetStore::enable_notifications() {
 95 |     notifications_enabled_ = true;
 96 | }
 97 | 
 98 | vector<string> OffsetStore::get_consumers() const {
 99 |     vector<string> output;
100 |     lock_guard<mutex> _(consumer_offsets_mutex_);
101 |     for (const auto& consumer_pair : consumer_offsets_) {
102 |         output.emplace_back(consumer_pair.first);
103 |     }
104 |     return output;
105 | }
106 | 
107 | vector<ConsumerOffset> OffsetStore::get_consumer_offsets(const string& group_id) const {
108 |     lock_guard<mutex> _(consumer_offsets_mutex_);
109 |     auto iter = consumer_offsets_.find(group_id);
110 |     if (iter == consumer_offsets_.end()) {
111 |         return {};
112 |     }
113 |     vector<ConsumerOffset> output;
114 |     for (const auto& topic_pair : iter->second) {
115 |         output.emplace_back(group_id, topic_pair.first.get_topic(),
116 |                             topic_pair.first.get_partition(), topic_pair.second);
117 |     }
118 |     return output;
119 | }
120 | 
121 | optional<int64_t> OffsetStore::get_topic_offset(const string& topic, int partition) const {
122 |     lock_guard<mutex> _(topic_offsets_mutex_);
123 |     auto iter = topic_offsets_.find({ topic, partition });
124 |     if (iter == topic_offsets_.end()) {
125 |         return boost::none;
126 |     }
127 |     return iter->second;
128 | }
129 | 
130 | vector<string> OffsetStore::get_topics() const {
131 |     return vector<string>(topics_.begin(), topics_.end());
132 | }
133 | 
134 | } // pirulo
135 | 


--------------------------------------------------------------------------------
/src/utils/task_scheduler.cpp:
--------------------------------------------------------------------------------
  1 | #include <algorithm>
  2 | #include "utils/task_scheduler.h"
  3 | #include "exceptions.h"
  4 | #include "detail/logging.h"
  5 | 
  6 | using std::lock_guard;
  7 | using std::unique_lock;
  8 | using std::mutex;
  9 | using std::thread;
 10 | using std::move;
 11 | using std::find_if;
 12 | using std::upper_bound;
 13 | using std::min;
 14 | using std::max;
 15 | 
 16 | using std::chrono::seconds;
 17 | 
 18 | namespace pirulo {
 19 | 
 20 | PIRULO_CREATE_LOGGER("p.scheduler");
 21 | 
 22 | static const double MINIMUM_PRIORITY = 0.1;
 23 | 
 24 | TaskScheduler::TaskScheduler() {
 25 |     process_thread_ = thread([&] {
 26 |         process_tasks();
 27 |     });
 28 | }
 29 | 
 30 | TaskScheduler::~TaskScheduler() {
 31 |     stop();
 32 | }
 33 | 
 34 | TaskScheduler::TaskId TaskScheduler::add_task(Task task, Duration maximum_offset) {
 35 |     lock_guard<mutex> _(tasks_mutex_);
 36 |     TaskId task_id = current_task_id_++;
 37 | 
 38 |     // Construct the new task and insert it
 39 |     TaskMetadata task_meta{ move(task), tasks_queue_.end(), ClockType::now(),
 40 |                             maximum_offset, 1.0 };
 41 |     auto iter = tasks_.emplace(task_id, move(task_meta)).first;
 42 | 
 43 |     // Schedule and update the stored iterator
 44 |     iter->second.queue_iterator = schedule_task(task_id, iter->second);
 45 |     return task_id;
 46 | }
 47 | 
 48 | void TaskScheduler::remove_task(TaskId id) {
 49 |     lock_guard<mutex> _(tasks_mutex_);
 50 |     auto iter = tasks_.find(id);
 51 |     if (iter == tasks_.end()) {
 52 |         return;
 53 |     }
 54 |     remove_scheduled_instance(id);
 55 |     tasks_.erase(iter);
 56 | }
 57 | 
 58 | void TaskScheduler::set_priority(TaskId id, double priority) {
 59 |     // Don't let this go too low
 60 |     priority = max(priority, MINIMUM_PRIORITY);
 61 | 
 62 |     lock_guard<mutex> _(tasks_mutex_);
 63 |     auto iter = tasks_.find(id);
 64 |     if (iter == tasks_.end()) {
 65 |         throw Exception("Task not found");
 66 |     }
 67 |     const auto now = ClockType::now();
 68 |     TaskMetadata& meta = iter->second; 
 69 |     const double priority_diff = priority - meta.priority;
 70 | 
 71 |     // Update the priority and the pariority set time to now
 72 |     meta.priority = priority;
 73 |     meta.last_priority_set_time = now;
 74 |     
 75 |     // Re-schedule this task if either:
 76 |     // * The prioity value is higher (meaning the priority is lower)
 77 |     // * The priority is lower but the task is already scheduled for more than our minimum
 78 |     // re-schedule time ahead
 79 |     const TaskExecutionInstance& current_instance = *meta.queue_iterator;
 80 |     if (priority_diff < 0 || now + minimum_reschedule_ < current_instance.scheduled_for) {
 81 |         remove_scheduled_instance(id);
 82 |         meta.queue_iterator = schedule_task(id, meta);
 83 |     }
 84 | }
 85 | 
 86 | void TaskScheduler::set_minimum_reschedule_time(Duration value) {
 87 |     minimum_reschedule_ = value;
 88 | }
 89 | 
 90 | void TaskScheduler::stop() {
 91 |     {
 92 |         lock_guard<mutex> _(tasks_mutex_);
 93 |         running_ = false;
 94 |         tasks_condition_.notify_all();
 95 |     }
 96 | 
 97 |     process_thread_.join();
 98 | }
 99 | 
100 | TaskScheduler::Duration TaskScheduler::get_schedule_delta(const TaskMetadata& meta) {
101 |     const size_t modifier = meta.maximum_offset.count() * meta.priority;;
102 |     return Duration(modifier);
103 | }
104 | 
105 | TaskScheduler::TaskQueue::iterator
106 | TaskScheduler::schedule_task(TaskId task_id, const TaskMetadata& meta) {
107 |     // Get the execution offset
108 |     const auto now = ClockType::now();
109 |     const auto schedule_for = now + get_schedule_delta(meta);
110 |     TaskExecutionInstance instance{task_id, now, schedule_for};
111 | 
112 |     // Find the task that's schedule after this one
113 |     const auto comparer = [](const TaskExecutionInstance& lhs, const TaskExecutionInstance& rhs) {
114 |         return lhs.scheduled_for < rhs.scheduled_for;
115 |     };
116 |     auto position = upper_bound(tasks_queue_.begin(), tasks_queue_.end(), instance, comparer);
117 |     return tasks_queue_.insert(position, move(instance));
118 | }
119 | 
120 | void TaskScheduler::remove_scheduled_instance(TaskId task_id) {
121 |     const auto comparer = [&](const TaskExecutionInstance& element) {
122 |         return element.task_id == task_id;
123 |     };
124 |     tasks_queue_.erase(find_if(tasks_queue_.begin(), tasks_queue_.end(), comparer));
125 | }
126 | 
127 | void TaskScheduler::process_tasks() {
128 |     while (running_) {
129 |         unique_lock<mutex> lock(tasks_mutex_);
130 |         auto now = ClockType::now();
131 |         while (running_) {
132 |             auto now = ClockType::now();
133 |             // Some random wake up time by default
134 |             auto wake_up_time = now + seconds(10);
135 |             if (!tasks_queue_.empty()) {
136 |                 wake_up_time = tasks_queue_.front().scheduled_for;
137 |             }
138 |             // If we should be awake already, stop loooping
139 |             if (wake_up_time <= now) {
140 |                 break;
141 |             }
142 |             tasks_condition_.wait_until(lock, wake_up_time);
143 |         }
144 |         // Make sure there's actually something to process
145 |         if (tasks_queue_.empty()) {
146 |             continue;
147 |         }
148 |         TaskExecutionInstance instance = move(tasks_queue_.front());
149 |         tasks_queue_.pop_front();
150 |         // Refresh the current wake_up_time
151 |         now = ClockType::now();
152 | 
153 |         // Before leaving the critical section, re-schedule the task
154 |         TaskMetadata& meta = tasks_.at(instance.task_id);
155 |         instance.scheduled_at = now;
156 |         instance.scheduled_for = now + get_schedule_delta(meta);
157 |         // If we're past our priority adjustment offset, this means the priority hasn't changed
158 |         // in a while. Update it accordingly
159 |         if (meta.last_priority_set_time + priority_adjustment_offset_ < now) {
160 |             const double new_priority = min(1.0, meta.priority * 2.0);
161 |             if (new_priority != meta.priority) {
162 |                 meta.priority = new_priority;
163 |                 LOG4CXX_TRACE(logger, "Set priority for task " << instance.task_id << " to "
164 |                               << meta.priority);
165 |             }
166 |         }
167 |         meta.queue_iterator = schedule_task(instance.task_id, meta);
168 | 
169 |         // Execute the task outside of the critical section
170 |         const Task task = meta.task;
171 |         lock.unlock();
172 |         task();
173 |     }
174 | }
175 | 
176 | } // pirulo
177 | 


--------------------------------------------------------------------------------
/include/detail/endianness.h:
--------------------------------------------------------------------------------
  1 | // Code taken from libtins: https://github.com/mfontanini/libtins/blob/master/include/tins/endianness.h
  2 | 
  3 | #pragma once
  4 | 
  5 | #include <stdint.h>
  6 | #if defined(__unix__) || (defined(__APPLE__) && defined(__MACH__))
  7 |     #include <sys/param.h>
  8 | #endif
  9 | 
 10 | #if defined(__APPLE__)
 11 |     #include <sys/types.h>
 12 |     #define PIRULO_IS_LITTLE_ENDIAN (BYTE_ORDER == LITTLE_ENDIAN)
 13 |     #define PIRULO_IS_BIG_ENDIAN (BYTE_ORDER == BIG_ENDIAN)
 14 | #elif defined(BSD)
 15 |     #include <sys/endian.h>
 16 |     #define PIRULO_IS_LITTLE_ENDIAN (_BYTE_ORDER == _LITTLE_ENDIAN)
 17 |     #define PIRULO_IS_BIG_ENDIAN (_BYTE_ORDER == _BIG_ENDIAN)
 18 | #elif defined(_WIN32)
 19 |     #include <cstdlib>
 20 |     #define PIRULO_IS_LITTLE_ENDIAN 1
 21 |     #define PIRULO_IS_BIG_ENDIAN 0
 22 | #else
 23 |     #include <endian.h>
 24 |     #define PIRULO_IS_LITTLE_ENDIAN (__BYTE_ORDER == __LITTLE_ENDIAN)
 25 |     #define PIRULO_IS_BIG_ENDIAN (__BYTE_ORDER == __BIG_ENDIAN)
 26 | #endif
 27 | 
 28 | // Define macros to swap bytes using compiler intrinsics when possible
 29 | #if defined(_MSC_VER)
 30 |     #define PIRULO_BYTE_SWAP_16(data) _byteswap_ushort(data)
 31 |     #define PIRULO_BYTE_SWAP_32(data) _byteswap_ulong(data)
 32 |     #define PIRULO_BYTE_SWAP_64(data) _byteswap_uint64(data)
 33 | #elif defined(PIRULO_HAVE_GCC_BUILTIN_SWAP)
 34 |     #define PIRULO_BYTE_SWAP_16(data) __builtin_bswap16(data)
 35 |     #define PIRULO_BYTE_SWAP_32(data) __builtin_bswap32(data)
 36 |     #define PIRULO_BYTE_SWAP_64(data) __builtin_bswap64(data)
 37 | #else
 38 |     #define PIRULO_NO_BYTE_SWAP_INTRINSICS
 39 | #endif
 40 | 
 41 | namespace pirulo {
 42 | namespace endian {
 43 | 
 44 | /** 
 45 |  * \brief "Changes" a 8-bit integral value's endianess. This is an
 46 |  * identity function.
 47 |  *
 48 |  * \param data The data to convert.
 49 |  */
 50 | inline uint8_t do_change_endian(uint8_t data) {
 51 |     return data;
 52 | }
 53 | 
 54 | /** 
 55 |  * \brief Changes a 16-bit integral value's endianess.
 56 |  *
 57 |  * \param data The data to convert.
 58 |  */
 59 | inline uint16_t do_change_endian(uint16_t data) {
 60 |     #ifdef PIRULO_NO_BYTE_SWAP_INTRINSICS
 61 |         return ((data & 0xff00) >> 8)  | ((data & 0x00ff) << 8);
 62 |     #else
 63 |         return PIRULO_BYTE_SWAP_16(data);
 64 |     #endif
 65 | }
 66 | 
 67 | /**
 68 |  * \brief Changes a 32-bit integral value's endianess.
 69 |  *
 70 |  * \param data The data to convert.
 71 |  */
 72 | inline uint32_t do_change_endian(uint32_t data) {
 73 |     #ifdef PIRULO_NO_BYTE_SWAP_INTRINSICS
 74 |         return (((data & 0xff000000) >> 24) | ((data & 0x00ff0000) >> 8)  |
 75 |                ((data & 0x0000ff00) << 8)  | ((data & 0x000000ff) << 24));
 76 |     #else
 77 |         return PIRULO_BYTE_SWAP_32(data);
 78 |     #endif
 79 | }
 80 | 
 81 | /**
 82 |  * \brief Changes a 64-bit integral value's endianess.
 83 |  *
 84 |  * \param data The data to convert.
 85 |  */
 86 |  inline uint64_t do_change_endian(uint64_t data) {
 87 |     #ifdef PIRULO_NO_BYTE_SWAP_INTRINSICS
 88 |         return (((uint64_t)(do_change_endian((uint32_t)(data & 0xffffffff))) << 32) |
 89 |                (do_change_endian(((uint32_t)(data >> 32)))));
 90 |     #else
 91 |         return PIRULO_BYTE_SWAP_64(data);
 92 |     #endif
 93 |  }
 94 |  
 95 | /**
 96 |  * \cond
 97 |  */
 98 | 
 99 | // Helpers to convert
100 | template<typename T>
101 | struct conversion_dispatch_helper {
102 |     static T dispatch(T data) {
103 |         return do_change_endian(data);
104 |     }
105 | };
106 | 
107 | 
108 | template<size_t>
109 | struct conversion_dispatcher;
110 | 
111 | template<>
112 | struct conversion_dispatcher<sizeof(uint8_t)> 
113 | : conversion_dispatch_helper<uint8_t> { };
114 | 
115 | template<>
116 | struct conversion_dispatcher<sizeof(uint16_t)> 
117 | : conversion_dispatch_helper<uint16_t> { };
118 | 
119 | template<>
120 | struct conversion_dispatcher<sizeof(uint32_t)> 
121 | : conversion_dispatch_helper<uint32_t> { };
122 | 
123 | template<>
124 | struct conversion_dispatcher<sizeof(uint64_t)> 
125 | : conversion_dispatch_helper<uint64_t> { };
126 | 
127 | /**
128 |  * \endcond
129 |  */
130 | 
131 | /**
132 |  * \brief Changes an integral value's endianess.
133 |  * 
134 |  * This dispatchs to the corresponding function.
135 |  *
136 |  * \param data The data to convert.
137 |  */ 
138 |  template<typename T>
139 |  inline T change_endian(T data) {
140 |      return conversion_dispatcher<sizeof(T)>::dispatch(data);
141 |  }
142 | 
143 | #if PIRULO_IS_LITTLE_ENDIAN
144 |     /** 
145 |      * \brief Convert any integral type to big endian.
146 |      *
147 |      * \param data The data to convert.
148 |      */
149 |     template<typename T>
150 |     inline T host_to_be(T data) {
151 |         return change_endian(data);
152 |     }
153 |      
154 |     /**
155 |      * \brief Convert any integral type to little endian.
156 |      *
157 |      * On little endian platforms, the parameter is simply returned.
158 |      * 
159 |      * \param data The data to convert.
160 |      */
161 |      template<typename T>
162 |      inline T host_to_le(T data) {
163 |          return data;
164 |      }
165 |      
166 |     /**
167 |      * \brief Convert any big endian value to the host's endianess.
168 |      * 
169 |      * \param data The data to convert.
170 |      */
171 |      template<typename T>
172 |      inline T be_to_host(T data) {
173 |          return change_endian(data);
174 |      }
175 |      
176 |     /**
177 |      * \brief Convert any little endian value to the host's endianess.
178 |      * 
179 |      * \param data The data to convert.
180 |      */
181 |      template<typename T>
182 |      inline T le_to_host(T data) {
183 |          return data;
184 |      }
185 | #elif PIRULO_IS_BIG_ENDIAN
186 |     /** 
187 |      * \brief Convert any integral type to big endian.
188 |      *
189 |      * \param data The data to convert.
190 |      */
191 |     template<typename T>
192 |     inline T host_to_be(T data) {
193 |         return data;
194 |     }
195 |      
196 |     /**
197 |      * \brief Convert any integral type to little endian.
198 |      *
199 |      * On little endian platforms, the parameter is simply returned.
200 |      * 
201 |      * \param data The data to convert.
202 |      */
203 |      template<typename T>
204 |      inline T host_to_le(T data) {
205 |          return change_endian(data);
206 |      }
207 |      
208 |     /**
209 |      * \brief Convert any big endian value to the host's endianess.
210 |      * 
211 |      * \param data The data to convert.
212 |      */
213 |      template<typename T>
214 |      inline T be_to_host(T data) {
215 |          return data;
216 |      }
217 |      
218 |     /**
219 |      * \brief Convert any little endian value to the host's endianess.
220 |      * 
221 |      * \param data The data to convert.
222 |      */
223 |      template<typename T>
224 |      inline T le_to_host(T data) {
225 |          return change_endian(data);
226 |      }
227 | #endif
228 |      
229 | } // endian
230 | } // pirulo
231 | 


--------------------------------------------------------------------------------
/src/topic_offset_reader.cpp:
--------------------------------------------------------------------------------
  1 | #include <tuple>
  2 | #include <cppkafka/metadata.h>
  3 | #include "topic_offset_reader.h"
  4 | #include "detail/logging.h"
  5 | #include "utils/utils.h"
  6 | 
  7 | using std::string;
  8 | using std::move;
  9 | using std::unordered_set;
 10 | using std::set;
 11 | using std::lock_guard;
 12 | using std::mutex;
 13 | using std::tie;
 14 | using std::ignore;
 15 | 
 16 | using std::this_thread::sleep_for;
 17 | 
 18 | using std::chrono::seconds;
 19 | 
 20 | using cppkafka::Consumer;
 21 | using cppkafka::Message;
 22 | using cppkafka::Configuration;
 23 | using cppkafka::Metadata;
 24 | using cppkafka::TopicMetadata;
 25 | using cppkafka::TopicPartition;
 26 | using cppkafka::TopicPartitionList;
 27 | 
 28 | namespace pirulo {
 29 | 
 30 | PIRULO_CREATE_LOGGER("p.topics");
 31 | 
 32 | static Configuration prepare_config(Configuration config) {
 33 |     config.set("group.id", utils::generate_group_id());
 34 |     LOG4CXX_INFO(logger, "Using consumer " << config.get("group.id") << " for "
 35 |                  << " topic offset consumption");
 36 |     return config;
 37 | }
 38 | 
 39 | TopicOffsetReader::TopicOffsetReader(StorePtr store, size_t thread_count,
 40 |                                      ConsumerOffsetReaderPtr consumer_reader,
 41 |                                      Configuration config)
 42 | : consumer_pool_(thread_count, prepare_config(move(config))),store_(move(store)),
 43 |   thread_pool_(thread_count), consumer_offset_reader_(move(consumer_reader)) {
 44 | 
 45 | }
 46 | 
 47 | void TopicOffsetReader::run() {
 48 |     LOG4CXX_INFO(logger, "Performing topic offsets cold start");
 49 |     process_metadata([&](TopicPartitionCount topics) {
 50 |         async_process_topics(topics);
 51 |         monitor_topics(topics);
 52 |     });
 53 |     thread_pool_.wait_for_tasks();
 54 |     LOG4CXX_INFO(logger, "Finished topic offsets cold start");
 55 | 
 56 |     monitor_new_topics();
 57 | }
 58 | 
 59 | void TopicOffsetReader::stop() {
 60 |     running_ = false;
 61 |     thread_pool_.stop();
 62 | }
 63 | 
 64 | TopicOffsetReader::StorePtr TopicOffsetReader::get_store() const {
 65 |     return store_;
 66 | }
 67 | 
 68 | void TopicOffsetReader::async_process_topics(const TopicPartitionCount& topics) {
 69 |     LOG4CXX_INFO(logger, "Fetching offsets for " << topics.size() << " topics");
 70 |     for (const auto& topic_pair : topics) {
 71 |         const string& topic = topic_pair.first;
 72 |         const size_t partition_count = topic_pair.second;
 73 |         for (size_t i = 0; i < partition_count; ++i) {
 74 |             thread_pool_.add_task([&, topic, i] {
 75 |                 process_topic_partition({ topic, static_cast<int>(i) });
 76 |             });
 77 |         }
 78 |     }
 79 | }
 80 | 
 81 | void TopicOffsetReader::monitor_topics(const TopicPartitionCount& topics) {
 82 |     // They're all new, just convert them to a TopicPartitionList
 83 |     monitor_topics(get_new_topic_partitions(topics));
 84 | }
 85 | 
 86 | void TopicOffsetReader::monitor_topics(const TopicPartitionList& topics) {
 87 |     for (const TopicPartition& topic_partition : topics) {
 88 |         auto task = [&, topic_partition] {
 89 |             process_topic_partition(topic_partition);
 90 |         };
 91 |         // Schedule a task to process it periodically
 92 |         auto task_id = task_scheduler_.add_task(move(task), maximum_topic_reload_time_);
 93 | 
 94 |         // Mark it as monitored
 95 |         monitored_topic_task_id_.emplace(topic_partition, task_id);
 96 | 
 97 |         // Watch for commits on this topic
 98 |         auto commit_callback = [&](const string& topic, int partition) {
 99 |             on_commit(topic, partition);
100 |         };
101 |         consumer_offset_reader_->watch_commits(topic_partition.get_topic(),
102 |                                                topic_partition.get_partition(),
103 |                                                move(commit_callback));
104 |     }
105 | }
106 | 
107 | void TopicOffsetReader::monitor_new_topics() {
108 |     auto task = [&] {
109 |         process_metadata([&](const TopicPartitionCount& counts) {
110 |             const TopicPartitionList new_topics = get_new_topic_partitions(counts);
111 |             if (!new_topics.empty()) {
112 |                 LOG4CXX_INFO(logger, "Found " << new_topics.size() << " new topic/partitions "
113 |                              "to process");
114 |                 monitor_topics(new_topics);
115 |             }
116 |         });
117 |     };
118 |     task_scheduler_.add_task(move(task), maximum_metadata_reload_time_);
119 | }
120 | 
121 | TopicPartitionList TopicOffsetReader::get_new_topic_partitions(const TopicPartitionCount& counts) {
122 |     TopicPartitionList output;
123 |     for (const auto& topic_count_pair : counts) {
124 |         const string& topic = topic_count_pair.first;
125 |         for (size_t i = 0; i < topic_count_pair.second; ++i) {
126 |             TopicPartition topic_partition(topic, i);
127 |             if (!monitored_topic_task_id_.count(topic_partition)) {
128 |                 output.emplace_back(move(topic_partition));
129 |             }
130 |         }
131 |     }
132 |     return output;
133 | }
134 | 
135 | TopicOffsetReader::TopicPartitionCount TopicOffsetReader::load_metadata() {
136 |     // Load all existing topic names
137 |     TopicPartitionCount topics;
138 |     consumer_pool_.acquire_consumer([&](Consumer& consumer) {
139 |         Metadata md = consumer.get_metadata();
140 |         for (const TopicMetadata& topic_metadata : md.get_topics()) {
141 |             topics.emplace(topic_metadata.get_name(),
142 |                            topic_metadata.get_partitions().size());
143 |         }
144 |     });
145 |     return topics;
146 | }
147 | 
148 | void TopicOffsetReader::process_metadata(const MetadataCallback& callback) {
149 |     LOG4CXX_INFO(logger, "Loading topics metadata");
150 | 
151 |     try {
152 |         TopicPartitionCount topics = load_metadata();
153 |         LOG4CXX_INFO(logger, "Finished loading metadata, found " << topics.size() << " topics");
154 |         callback(move(topics));
155 |     }
156 |     catch (const cppkafka::Exception& ex) {
157 |         LOG4CXX_ERROR(logger, "Failed to fetch topic metadata: " << ex.what());
158 |     }
159 | }
160 | 
161 | void TopicOffsetReader::process_topic_partition(const TopicPartition& topic_partition) {
162 |     LOG4CXX_TRACE(logger, "Fetching offset for " << topic_partition);
163 |     uint64_t offset;
164 |     try {
165 |         consumer_pool_.acquire_consumer([&](Consumer& consumer) {
166 |             tie(ignore, offset) = consumer.query_offsets(topic_partition);
167 |             store_->store_topic_offset(topic_partition.get_topic(),
168 |                                        topic_partition.get_partition(), offset);
169 |         });
170 |     }
171 |     catch (const cppkafka::Exception& ex) {
172 |         LOG4CXX_ERROR(logger, "Failed to fetch offsets for " << topic_partition
173 |                       << ": " << ex.what());
174 |     }
175 | }
176 | 
177 | void TopicOffsetReader::on_commit(const string& topic, int partition) {
178 |     TopicPartition topic_partition(topic, partition);
179 |     LOG4CXX_TRACE(logger, "Bumping up priority of offset loading for " << topic_partition);
180 |     // Increase the priority for this task
181 |     auto task_id = monitored_topic_task_id_.at(topic_partition);
182 |     task_scheduler_.set_priority(task_id, 0.0);
183 | }
184 | 
185 | } // pirulo
186 | 


--------------------------------------------------------------------------------
/src/python/api.cpp:
--------------------------------------------------------------------------------
  1 | #include <boost/optional.hpp>
  2 | #include <boost/python/module.hpp>
  3 | #include <boost/python/class.hpp>
  4 | #include <boost/python/object.hpp>
  5 | #include <boost/python/str.hpp>
  6 | #include <boost/python/import.hpp>
  7 | #include <boost/python/wrapper.hpp>
  8 | #include <boost/python/call_method.hpp>
  9 | #include <boost/python/suite/indexing/vector_indexing_suite.hpp>
 10 | #include "python/api.h"
 11 | #include "python/helpers.h"
 12 | #include "python/handler.h"
 13 | #include "python/lag_tracker_handler.h"
 14 | #include "detail/logging.h"
 15 | 
 16 | using std::string;
 17 | using std::vector;
 18 | using std::bind;
 19 | using std::ref;
 20 | using std::shared_ptr;
 21 | 
 22 | using boost::optional;
 23 | 
 24 | namespace python = boost::python;
 25 | 
 26 | namespace pirulo {
 27 | namespace api {
 28 | 
 29 | PIRULO_CREATE_LOGGER("p.python");
 30 | 
 31 | template <typename Method, typename... Args>
 32 | void exec_method(const Method& functor, const Args&... args) {
 33 |     if (functor) {
 34 |         helpers::safe_exec(logger, bind(functor, ref(args)...));
 35 |     }
 36 |     else {
 37 |         LOG4CXX_DEBUG(logger, "Not executing callback as it's None");
 38 |     }
 39 | }
 40 | 
 41 | template <typename... Args>
 42 | void exec_method(python::object& object, const char* name, const Args&... args) {
 43 |     if (PyObject_HasAttrString(object.ptr(), name)) {
 44 |         const auto functor = object.attr(name);
 45 |         helpers::safe_exec(logger, bind(functor, ref(args)...)); 
 46 |     }
 47 | }
 48 | 
 49 | class HandlerWrapper : public Handler,
 50 |                        public python::wrapper<Handler> {
 51 | private:
 52 |     void handle_initialize() {
 53 |         exec_method(get_override("handle_initialize"));
 54 |     }
 55 | 
 56 |     void handle_new_consumer(const string& group_id) {
 57 |         exec_method(get_override("handle_new_consumer"), group_id);
 58 |     }
 59 | 
 60 |     void handle_new_topic(const string& topic) {
 61 |         exec_method(get_override("handle_new_topic"), topic);
 62 |     }
 63 | 
 64 |     void handle_consumer_commit(const string& group_id, const string& topic,
 65 |                                 int partition, int64_t offset) {
 66 |         exec_method(get_override("handle_consumer_commit"), group_id, topic, partition, offset);
 67 |     }
 68 | 
 69 |     void handle_topic_message(const string& topic, int partition, int64_t offset) {
 70 |         exec_method(get_override("handle_topic_message"), topic, partition, offset);
 71 |     }
 72 | };
 73 | 
 74 | class LagTrackerHandlerWrapper : public LagTrackerHandler,
 75 |                                  public python::wrapper<LagTrackerHandler> {
 76 | public:
 77 |     LagTrackerHandlerWrapper(PyObject* self)
 78 |     : self_(python::handle<>(self)) {
 79 | 
 80 |     }
 81 | private:
 82 |     void handle_initialize() {
 83 |         LagTrackerHandler::handle_initialize();
 84 |         exec_method(self_, "handle_initialize");
 85 |     }
 86 | 
 87 |     void handle_lag_update(const string& topic, int partition,
 88 |                            const string& group_id, uint64_t consumer_lag) {
 89 |         exec_method(self_, "handle_lag_update", topic, partition, group_id, consumer_lag);
 90 |     }
 91 | 
 92 |     python::object self_;
 93 | };
 94 | 
 95 | BOOST_PYTHON_MODULE(pirulo) {
 96 |     using python::class_;
 97 |     using python::bases;
 98 |     using python::init;
 99 |     using python::no_init;
100 |     using python::return_internal_reference;
101 | 
102 |     class_<HandlerWrapper, boost::noncopyable>("Handler")
103 |         .def("initialize", &Handler::initialize)
104 |         .def("get_offset_store", &Handler::get_offset_store, return_internal_reference<>())
105 |         .def("subscribe_to_consumers", &Handler::subscribe_to_consumers)
106 |         .def("subscribe_to_consumer_commits", &Handler::subscribe_to_consumer_commits)
107 |         .def("subscribe_to_topics", &Handler::subscribe_to_topics)
108 |         .def("subscribe_to_topic_message", &Handler::subscribe_to_topic_message)
109 |     ;
110 | 
111 |     class_<LagTrackerHandlerWrapper, bases<Handler>, LagTrackerHandlerWrapper,
112 |            boost::noncopyable>("LagTrackerHandler")
113 | 
114 |     ;
115 | }
116 | 
117 | template <typename T>
118 | struct value_or_none {
119 |     static PyObject* convert(const optional<T>& value) {
120 |         if (value) {
121 |             return python::incref(python::object(*value).ptr());
122 |         }
123 |         else {
124 |             return Py_None;
125 |         }
126 |     }
127 | };
128 | 
129 | void register_module() {
130 |     PyImport_AppendInittab("pirulo", &initpirulo);
131 | }
132 | 
133 | void register_types() {
134 |      using python::to_python_converter;
135 |      using python::class_;
136 |      using python::no_init;
137 |      using python::make_function;
138 |      using python::return_internal_reference;
139 |      using python::vector_indexing_suite;
140 |      using python::object;
141 |      using python::call;
142 | 
143 |     to_python_converter<optional<int64_t>, value_or_none<int64_t>>();
144 | 
145 |     class_<ConsumerOffset>("ConsumerOffset", no_init)
146 |         .add_property("group_id",
147 |                       make_function(&ConsumerOffset::get_group_id,
148 |                                     return_internal_reference<>()))
149 |         .add_property("topic", +[](const ConsumerOffset& o) {
150 |             return o.get_topic_partition().get_topic();
151 |         })
152 |         .add_property("partition", +[](const ConsumerOffset& o) {
153 |             return o.get_topic_partition().get_partition();
154 |         })
155 |         .add_property("offset", +[](const ConsumerOffset& o) {
156 |             return o.get_topic_partition().get_offset();
157 |         })
158 |         ;
159 | 
160 |     class_<OffsetStore, shared_ptr<OffsetStore>, boost::noncopyable>("OffsetStore", no_init)
161 |         .def("get_consumers", &OffsetStore::get_consumers)
162 |         .def("get_consumer_offsets", &OffsetStore::get_consumer_offsets)
163 |         .def("get_topic_offset", &OffsetStore::get_topic_offset)
164 |         .def("get_topics", &OffsetStore::get_topics)
165 |         .def("on_new_consumer", +[](OffsetStore& store, const object& callback) {
166 |             store.on_new_consumer([=](const string& group_id) {
167 |                 helpers::safe_exec(logger, [&]() {
168 |                     call<void>(callback.ptr(), group_id);
169 |                 });
170 |             });
171 |         })
172 |         .def("on_new_topic", +[](OffsetStore& store, const object& callback) {
173 |             store.on_new_topic([=](const string& topic) {
174 |                 helpers::safe_exec(logger, [&]() {
175 |                     call<void>(callback.ptr(), topic);
176 |                 });
177 |             });
178 |         })
179 |         .def("on_consumer_commit", +[](OffsetStore& store, const string& group_id,
180 |                                        const object& callback) {
181 |             store.on_consumer_commit(group_id, [=](const string& group_id, const string& topic,
182 |                                                    int partition, uint64_t offset) {
183 |                 helpers::safe_exec(logger, [&]() {
184 |                     call<void>(callback.ptr(), group_id, topic, partition, offset);
185 |                 });
186 |             });
187 |         })
188 |         .def("on_topic_message", +[](OffsetStore& store, const string& topic,
189 |                                      const object& callback) {
190 |             store.on_topic_message(topic, [=](const string& topic, int partition,
191 |                                               uint64_t offset) {
192 |                 helpers::safe_exec(logger, [&]() {
193 |                     call<void>(callback.ptr(), topic, partition, offset);
194 |                 });
195 |             });
196 |         })
197 |         ;
198 | 
199 |     class_<vector<string>>("StringVector")
200 |         .def(vector_indexing_suite<vector<string>>())
201 |         ;
202 | 
203 |     class_<vector<ConsumerOffset>>("ConsumerOffsetVector")
204 |         .def(vector_indexing_suite<vector<ConsumerOffset>>())
205 |         ;
206 | }
207 | 
208 | } // api
209 | } // pirulo
210 | 


--------------------------------------------------------------------------------