kafka_consumer.hpp

/*
 * Copyright (c) 2020-2022, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * https://apache.ac.cn/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once

#include "kafka_callback.hpp"

#include <cudf/io/datasource.hpp>

#include <librdkafka/rdkafkacpp.h>

#include <algorithm>
#include <chrono>
#include <map>
#include <memory>
#include <string>

namespace cudf {
namespace io {
namespace external {
namespace kafka {

class kafka_consumer : public cudf::io::datasource {
 public
 kafka_consumer(std::map<std::string, std::string> configs,
 python_callable_type python_callable,
 kafka_oauth_callback_wrapper_type callable_wrapper);

 kafka_consumer(std::map<std::string, std::string> configs,
 python_callable_type python_callable,
 kafka_oauth_callback_wrapper_type callable_wrapper,
 std::string const& topic_name,
 int partition,
 int64_t start_offset,
 int64_t end_offset,
 int batch_timeout,
 std::string const& delimiter);

 std::unique_ptr<cudf::io::datasource::buffer> host_read(size_t offset, size_t size) override;

 size_t size() const override;

 size_t host_read(size_t offset, size_t size, uint8_t* dst) override;

 void commit_offset(std::string const& topic, int partition, int64_t offset);

 std::map<std::string, int64_t> get_watermark_offset(std::string const& topic,
 int partition,
 int timeout,
 bool cached);

 std::map<std::string, std::string> current_configs();

 int64_t get_committed_offset(std::string const& topic, int partition);

 std::map<std::string, std::vector<int32_t>> list_topics(std::string specific_topic);

 void close(int timeout);

 void unsubscribe();

 virtual ~kafka_consumer(){};

 private
 std::unique_ptr<RdKafka::Conf> kafka_conf; // RDKafka 配置对象
 std::unique_ptr<RdKafka::KafkaConsumer> consumer;

 std::map<std::string, std::string> configs;
 python_callable_type python_callable_;
 kafka_oauth_callback_wrapper_type callable_wrapper_;

 std::string topic_name;
 int partition;
 int64_t start_offset;
 int64_t end_offset;
 int batch_timeout;
 int default_timeout = 10000; // 毫秒
 std::string delimiter;

 std::string buffer;

 private
 RdKafka::ErrorCode update_consumer_topic_partition_assignment(std::string const& topic,
 int partition,
 int64_t offset);

 int64_t now();

 void consume_to_buffer();
};

} // namespace kafka
} // namespace external
} // namespace io
} // namespace cudf