From 3220970fea30535d341f2735333518d57e1a0f27 Mon Sep 17 00:00:00 2001
From: Marco Neumann <marco@crepererum.net>
Date: Fri, 12 Nov 2021 11:31:45 +0100
Subject: [PATCH] fix: server IDs for routers on startup

---
 influxdb_iox/src/commands/run/router.rs       |  8 ++++
 .../tests/end_to_end_cases/deployment_api.rs  | 40 ++++++++++++++++++-
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/influxdb_iox/src/commands/run/router.rs b/influxdb_iox/src/commands/run/router.rs
index b0066976c6..087a56ff3a 100644
--- a/influxdb_iox/src/commands/run/router.rs
+++ b/influxdb_iox/src/commands/run/router.rs
@@ -12,6 +12,7 @@ use crate::{
     },
     structopt_blocks::run_config::RunConfig,
 };
+use observability_deps::tracing::warn;
 use router::{resolver::RemoteTemplate, server::RouterServer};
 use structopt::StructOpt;
 use thiserror::Error;
@@ -73,6 +74,13 @@ pub async fn command(config: Config) -> Result<()> {
         )
         .await,
     );
+    if let Some(id) = config.run_config.server_id_config.server_id {
+        router_server
+            .set_server_id(id)
+            .expect("server id already set");
+    } else {
+        warn!("server ID not set. ID must be set via the INFLUXDB_IOX_ID config or API before writing or querying data.");
+    }
     let server_type = Arc::new(RouterServerType::new(router_server, &common_state));
 
     Ok(influxdb_ioxd::main(common_state, server_type).await?)
diff --git a/influxdb_iox/tests/end_to_end_cases/deployment_api.rs b/influxdb_iox/tests/end_to_end_cases/deployment_api.rs
index 1bae9d22fa..0acfc8b798 100644
--- a/influxdb_iox/tests/end_to_end_cases/deployment_api.rs
+++ b/influxdb_iox/tests/end_to_end_cases/deployment_api.rs
@@ -7,7 +7,7 @@ use influxdb_iox_client::{
 use test_helpers::assert_error;
 use tonic::Code;
 
-use crate::common::server_fixture::{ServerFixture, ServerType};
+use crate::common::server_fixture::{ServerFixture, ServerType, TestConfig};
 
 #[tokio::test]
 async fn test_serving_readiness_database() {
@@ -104,6 +104,28 @@ async fn test_set_get_server_id_router() {
     assert_set_get_server_id(ServerFixture::create_single_use(ServerType::Router).await).await;
 }
 
+#[tokio::test]
+async fn test_set_get_server_id_already_set_database() {
+    let server_id = NonZeroU32::new(42).unwrap();
+    let test_config = TestConfig::new(ServerType::Database).with_server_id(server_id);
+    assert_set_get_server_id_already_set(
+        ServerFixture::create_single_use_with_config(test_config).await,
+        server_id,
+    )
+    .await;
+}
+
+#[tokio::test]
+async fn test_set_get_server_id_already_set_router() {
+    let server_id = NonZeroU32::new(42).unwrap();
+    let test_config = TestConfig::new(ServerType::Router).with_server_id(server_id);
+    assert_set_get_server_id_already_set(
+        ServerFixture::create_single_use_with_config(test_config).await,
+        server_id,
+    )
+    .await;
+}
+
 async fn assert_set_get_server_id(server_fixture: ServerFixture) {
     let mut client = server_fixture.deployment_client();
 
@@ -123,3 +145,19 @@ async fn assert_set_get_server_id(server_fixture: ServerFixture) {
         .await;
     assert_error!(result, UpdateServerIdError::AlreadySet);
 }
+
+async fn assert_set_get_server_id_already_set(
+    server_fixture: ServerFixture,
+    server_id: NonZeroU32,
+) {
+    let mut client = server_fixture.deployment_client();
+
+    let got = client.get_server_id().await.expect("get ID failed");
+    assert_eq!(got, server_id);
+
+    // setting server ID a second time should fail
+    let result = client
+        .update_server_id(NonZeroU32::try_from(13).unwrap())
+        .await;
+    assert_error!(result, UpdateServerIdError::AlreadySet);
+}