core, osrdyne, docker: make rmq max message size configurable through an environment variable

Signed-off-by: Ethan Perruzza <[email protected]>

Author: EthanPERRUZZA

core/src/main/java/fr/sncf/osrd/cli/WorkerCommand.kt

@@ -45,6 +45,7 @@ class WorkerCommand : CliCommand {
     val WORKER_ID_USE_HOSTNAME: Boolean
     val WORKER_KEY: String?
     val WORKER_AMQP_URI: String
+    val WORKER_MAX_MSG_SIZE: Int
     val WORKER_POOL: String
     val WORKER_REQUESTS_QUEUE: String
     val WORKER_ACTIVITY_EXCHANGE: String
@@ -56,6 +57,7 @@ class WorkerCommand : CliCommand {
         WORKER_KEY = if (ALL_INFRA) "all" else System.getenv("WORKER_KEY")
         WORKER_AMQP_URI =
             System.getenv("WORKER_AMQP_URI") ?: "amqp://osrd:[email protected]:5672/%2f"
+        WORKER_MAX_MSG_SIZE = getIntEnvvar("WORKER_MAX_MSG_SIZE") ?: 1024 * 1024 * 128 * 5
         WORKER_POOL = System.getenv("WORKER_POOL") ?: "core"
         WORKER_REQUESTS_QUEUE =
             System.getenv("WORKER_REQUESTS_QUEUE") ?: "$WORKER_POOL-req-$WORKER_KEY"
@@ -76,6 +78,10 @@ class WorkerCommand : CliCommand {
         return System.getenv(name)?.lowercase() !in arrayOf(null, "", "0", "false")
     }
 
+    private fun getIntEnvvar(name: String): Int? {
+        return System.getenv(name)?.toIntOrNull()
+    }
+
     override fun run(): Int {
         if (WORKER_ID == null || WORKER_KEY == null) {
             throw IllegalStateException(
@@ -125,7 +131,7 @@ class WorkerCommand : CliCommand {
 
         val factory = ConnectionFactory()
         factory.setUri(WORKER_AMQP_URI)
-        factory.setMaxInboundMessageBodySize(1024 * 1024 * 128 * 5)
+        factory.setMaxInboundMessageBodySize(WORKER_MAX_MSG_SIZE)
         val connection = factory.newConnection()
         connection.createChannel().use { channel -> reportActivity(channel, "started") }
 

docker/docker-compose.host.yml

@@ -54,6 +54,7 @@ services:
       OSRDYNE__WORKER_DRIVER__DEFAULT_ENV: >
         [ "CORE_EDITOAST_URL=http://127.0.0.1:8090", "JAVA_TOOL_OPTIONS=-javaagent:/app/opentelemetry-javaagent.jar", "CORE_MONITOR_TYPE=opentelemetry", "OTEL_EXPORTER_OTLP_TRACES_PROTOCOL=grpc", "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT=http://127.0.0.1:4317", "OTEL_METRICS_EXPORTER=none", "OTEL_LOGS_EXPORTER=none" ]
       OSRDYNE__AMQP_URI: "amqp://osrd:[email protected]:5672/%2f"
+      OSRDYNE__MAX_MSG_SIZE: 671088640 # 1024 * 1024 * 128 * 5
       OSRDYNE__MANAGEMENT_URI: "http://osrd:[email protected]:15672/"
       OSRDYNE__OPENTELEMETRY__ENDPOINT: "http://localhost:4317"
 

osrdyne/configuration_templates/process-compose-driver.yaml

@@ -9,3 +9,4 @@ worker_driver:
     worker_id: /tmp/worker_id # WORKER_ID
     worker_key: /tmp/worker_key # WORKER_KEY
     amqp_uri: /tmp/worker_amqp_uri # WORKER_AMQP_URI
+    max_msg_size: /tmp/worker_max_msg_size # WORKER_MAX_MSG_SIZE

osrdyne/configuration_templates/process-compose.yaml

@@ -19,6 +19,7 @@ processes:
         export WORKER_ID=$(cat /tmp/worker_id);
         export WORKER_KEY=$(cat /tmp/worker_key);
         export WORKER_AMQP_URI="$(cat /tmp/worker_amqp_uri);"
+        export WORKER_MAX_MSG_SIZE="$(cat /tmp/worker_max_msg_size);"
         java $$JAVA_OPTS -ea -jar build/libs/osrd-all.jar worker
       )
     working_dir: "$OSRD_PATH/core"

osrdyne/src/config.rs

@@ -25,6 +25,7 @@ pub enum WorkerDriverConfig {
 #[derive(Debug, Deserialize, Serialize)]
 pub struct OsrdyneConfig {
     pub amqp_uri: String,
+    pub max_msg_size: i64,
     pub management_uri: String,
     pub pool_id: String,
     pub worker_driver: WorkerDriverConfig,
@@ -56,6 +57,7 @@ impl Default for OsrdyneConfig {
     fn default() -> Self {
         Self {
             amqp_uri: "amqp://osrd:password@osrd-rabbitmq:5672/%2f".into(),
+            max_msg_size: 1024 * 1024 * 128 * 5,
             management_uri: "http://osrd:password@osrd-rabbitmq:15672".into(),
             pool_id: "core".to_string(),
             worker_driver: WorkerDriverConfig::Noop,

osrdyne/src/drivers/docker.rs

@@ -41,6 +41,7 @@ pub struct DockerDriver {
     client: Docker,
     options: DockerDriverOptions,
     amqp_uri: String,
+    max_message_size: i64,
     worker_pool: String,
     version_identifier: String,
 }
@@ -60,6 +61,7 @@ impl DockerDriver {
     pub fn new(
         options: DockerDriverOptions,
         amqp_uri: String,
+        max_message_size: i64,
         worker_pool: String,
     ) -> DockerDriver {
         let version_identifier = std::env::var("OSRD_GIT_DESCRIBE")
@@ -71,6 +73,7 @@ impl DockerDriver {
             client: Docker::connect_with_socket_defaults().expect("Failed to connect to Docker"),
             options,
             amqp_uri,
+            max_message_size,
             worker_pool,
             version_identifier: hashed,
         }
@@ -120,6 +123,7 @@ impl WorkerDriver for DockerDriver {
                 env.push(format!("WORKER_ID={}", new_id));
                 env.push(format!("WORKER_KEY={}", worker_key));
                 env.push(format!("WORKER_AMQP_URI={}", self.amqp_uri));
+                env.push(format!("WORKER_MAX_MSG_SIZE={}", self.max_message_size));
                 env
             };
 

osrdyne/src/drivers/kubernetes.rs

@@ -139,6 +139,7 @@ pub struct KubernetesDriver {
     pool_id: String,
     options: KubernetesDriverOptions,
     amqp_uri: String,
+    max_message_size: i64,
     version_identifier: String,
 }
 
@@ -157,6 +158,7 @@ impl KubernetesDriver {
     pub async fn new(
         options: KubernetesDriverOptions,
         amqp_uri: String,
+        max_message_size: i64,
         pool_id: String,
     ) -> KubernetesDriver {
         let version_identifier = std::env::var("OSRD_GIT_DESCRIBE")
@@ -170,6 +172,7 @@ impl KubernetesDriver {
                 .expect("Failed to connect to Kubernetes"),
             options,
             amqp_uri,
+            max_message_size,
             pool_id,
             version_identifier: hashed,
         }
@@ -408,6 +411,11 @@ impl WorkerDriver for KubernetesDriver {
                     value: Some(self.amqp_uri.clone()),
                     ..Default::default()
                 });
+                env.push(EnvVar {
+                    name: "WORKER_MAX_MSG_SIZE".to_string(),
+                    value: Some(self.max_message_size.to_string()),
+                    ..Default::default()
+                });
                 env.push(EnvVar {
                     name: "WORKEY_ID_USE_HOSTNAME".to_string(),
                     value: Some("1".to_string()),

osrdyne/src/main.rs

@@ -131,15 +131,21 @@ async fn main() -> Result<(), anyhow::Error> {
                 Box::new(DockerDriver::new(
                     opts,
                     config.amqp_uri.clone(),
+                    config.max_msg_size,
                     config.pool_id.clone(),
                 ))
             }
 
             WorkerDriverConfig::KubernetesDriver(opts) => {
                 info!("Using Kubernetes driver");
                 Box::new(
-                    KubernetesDriver::new(opts, config.amqp_uri.clone(), config.pool_id.clone())
-                        .await,
+                    KubernetesDriver::new(
+                        opts,
+                        config.amqp_uri.clone(),
+                        config.max_msg_size,
+                        config.pool_id.clone(),
+                    )
+                    .await,
                 )
             }