mirror of
https://github.com/thilo-behnke/wasm-pong.git
synced 2026-07-11 02:29:25 +00:00
working implementation for add partition
This commit is contained in:
@@ -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" ]
|
||||
|
||||
@@ -1,2 +1 @@
|
||||
/.idea
|
||||
/target
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user