working implementation for add partition

This commit is contained in:
Thilo Behnke
2022-05-28 19:17:25 +02:00
parent 0283e0da76
commit bf74e09cae
3 changed files with 70 additions and 24 deletions

View File

@@ -11,7 +11,7 @@ ADD custom-entrypoint.sh /
COPY --from=build-stage /opt/pong/kafka-script-proxy/target/release/kafka-script-proxy /bin/
USER root
RUN mkdir -p /var/lib/kafka-script-proxy
RUN mkdir -p /var/log/kafka-script-proxy
ENTRYPOINT []
CMD [ "/custom-entrypoint.sh" ]

View File

@@ -1,2 +1 @@
/.idea
/target

View File

@@ -1,8 +1,8 @@
use hyper::server::conn::AddrStream;
extern crate core;
use hyper::service::{make_service_fn, service_fn};
use hyper::{Body, Method, Request, Response, Server};
use hyper::{Body, Method, Request, Response, Server, StatusCode};
use std::convert::Infallible;
use std::fs::write;
use std::process::Output;
use tokio::fs::OpenOptions;
use tokio::io::AsyncWriteExt;
@@ -16,7 +16,7 @@ pub async fn main() {
}
pub async fn run() -> Result<(), ()> {
let make_svc = make_service_fn(|socket: &AddrStream| async {
let make_svc = make_service_fn(|_| async {
Ok::<_, Infallible>(service_fn(|req: Request<Body>| async {
handle_request(req).await
}))
@@ -26,7 +26,7 @@ pub async fn run() -> Result<(), ()> {
let server = Server::bind(&host).serve(make_svc);
write_to_log(&format!("Listening on http://{}", host)).await;
let graceful = server.with_graceful_shutdown(shutdown_signal());
graceful.await;
graceful.await.unwrap();
Ok(())
}
@@ -39,47 +39,91 @@ async fn handle_request(req: Request<Body>) -> Result<Response<Body>, Infallible
.await;
match (req.method(), req.uri().path()) {
(&Method::POST, "/add_partition") => handle_add_partition().await,
_ => Ok(Response::new("unknown".into())),
_ => build_error_res("not found", StatusCode::NOT_FOUND),
}
}
async fn handle_add_partition() -> Result<Response<Body>, Infallible> {
write_to_log("Called to add partition.").await;
let current_count = get_highest_partition_count().await;
if let Err(_) = current_count {
let err = "Failed to retrieve max partition count.";
write_to_log(err).await;
return build_error_res(err, StatusCode::INTERNAL_SERVER_ERROR);
}
let current_count = current_count.unwrap();
write_to_log(&format!("Successfully retrieved current max partition count: {}.", current_count)).await;
let next_partition = current_count + 1;
write_to_log(&format!("Updating partition count to {} for the following topics: {}", next_partition, TOPICS.join(","))).await;
for topic in TOPICS {
let output = run_command(&format!("/opt/bitnami/kafka/bin/kafka-topics.sh --bootstrap-server localhost:9092 --alter --topic {} --partitions {}", topic, next_partition)).await;
if let Err(e) = output {
write_to_log(&format!("{}", e)).await;
return Ok(Response::new(Body::empty()));
let error = format!("Failed to update the partition count: {}", e);
write_to_log(&error).await;
return build_error_res(&error, StatusCode::INTERNAL_SERVER_ERROR);
}
let output = output.unwrap();
if !output.status.success() {
write_to_log(&format!("{:?}", std::str::from_utf8(&*output.stderr))).await;
return Ok(Response::new(Body::empty()));
let error = format!("Failed to update the partition count: {:?}", std::str::from_utf8(&*output.stderr).unwrap());
write_to_log(&error).await;
return build_error_res(&error, StatusCode::INTERNAL_SERVER_ERROR);
}
}
Ok(Response::new(Body::from(next_partition.to_string())))
write_to_log(&format!("Successfully updated partition count to {}", next_partition)).await;
return build_success_res(&next_partition.to_string());
}
async fn get_highest_partition_count() -> u32 {
async fn get_highest_partition_count() -> Result<u32, PartitionCountQueryError> {
let output = run_command("/opt/bitnami/kafka/bin/kafka-topics.sh --bootstrap-server localhost:9092 --describe | grep -Po '(?<=PartitionCount: )(\\d+)' | sort -r | head -1").await.unwrap();
let stdout = std::str::from_utf8(&*output.stdout).unwrap().to_owned();
stdout.parse::<u32>().unwrap()
let output_str = std::str::from_utf8(&*output.stdout);
if let Err(e) = output_str {
let message = format!("Failed to convert command output to string: {:?}", e);
write_to_log(&message).await;
return Err(PartitionCountQueryError { message });
}
let output_str = output_str.unwrap().trim().replace("\n", "");
let parse_res = output_str.parse::<u32>();
if let Err(e) = parse_res {
let message = format!("Failed to parse partition count for output {}: {:?}", output_str, e);
write_to_log(&message).await;
return Err(PartitionCountQueryError {message})
}
return Ok(parse_res.unwrap())
}
#[derive(Debug)]
struct PartitionCountQueryError {
message: String
}
async fn run_command(command: &str) -> std::io::Result<Output> {
write_to_log(&format!("Running command: {}", command)).await;
let output = Command::new("/bin/bash")
.arg(format!("-c {}", command))
.arg("-c")
.arg(command)
.output()
.await.unwrap();
let stdout = std::str::from_utf8(&*output.stdout).unwrap();
let stderr = std::str::from_utf8(&*output.stderr).unwrap();
write_to_log(&format!("Command returned stdout: {}", stdout));
write_to_log(&format!("Command returned stdout: {}", stderr));
write_to_log(&format!("Command returned stdout: {}", stdout)).await;
write_to_log(&format!("Command returned stderr: {}", stderr)).await;
Ok(output)
}
pub fn build_success_res(value: &str) -> Result<Response<Body>, Infallible> {
let json = format!("{{\"data\": {}}}", value);
return Ok(Response::new(Body::from(json)));
}
pub fn build_error_res(error: &str, status: StatusCode) -> Result<Response<Body>, Infallible> {
let json = format!("{{\"error\": {}}}", error);
let mut res = Response::new(Body::from(json));
*res.status_mut() = status;
return Ok(res);
}
pub async fn shutdown_signal() {
// Wait for the CTRL+C signal
tokio::signal::ctrl_c()
@@ -87,18 +131,21 @@ pub async fn shutdown_signal() {
.expect("failed to install CTRL+C signal handler");
}
async fn write_to_log(value: &str) -> std::io::Result<()> {
async fn write_to_log(value: &str) {
let content = format!("{}\n", value);
let mut options = OpenOptions::new();
let mut file = options
let file = options
.create(true)
.write(true)
.append(true)
.open("/var/lib/kafka-script-proxy/log.txt")
.open("/var/log/kafka-script-proxy/log.txt")
.await;
if let Err(e) = file {
println!("Failed to open file: {}", e);
return Ok(());
return;
}
let write_res = file.unwrap().write_all(content.as_bytes()).await;
if let Err(e) = write_res {
println!("Failed to write to file: {:?}", e);
}
file.unwrap().write_all(content.as_bytes()).await
}