diff --git a/src/control-plane-services/instance-cluster-management/NOTICE b/src/control-plane-services/instance-cluster-management/NOTICE index 6b0a8a11e..bacab7ca2 100644 --- a/src/control-plane-services/instance-cluster-management/NOTICE +++ b/src/control-plane-services/instance-cluster-management/NOTICE @@ -1,5 +1,5 @@ -Lists of 293 third-party dependencies. +Lists of 290 third-party dependencies. (Apache License, Version 2.0) LZ4 Java Compression (at.yawk.lz4:lz4-java:1.10.3 - https://github.com/yawkat/lz4-java) (EPL-2.0) (LGPL-2.1-only) Logback Classic Module (ch.qos.logback:logback-classic:1.5.37 - http://logback.qos.ch) (EPL-2.0) (LGPL-2.1-only) Logback Core Module (ch.qos.logback:logback-core:1.5.37 - http://logback.qos.ch) @@ -69,9 +69,6 @@ Lists of 293 third-party dependencies. (The Apache Software License, Version 2.0) micrometer-tracing (io.micrometer:micrometer-tracing:1.6.6 - https://github.com/micrometer-metrics/tracing) (The Apache Software License, Version 2.0) micrometer-tracing-bridge-otel (io.micrometer:micrometer-tracing-bridge-otel:1.6.6 - https://github.com/micrometer-metrics/tracing) (The Apache License, Version 2.0) jnats (io.nats:jnats:2.23.0 - https://github.com/nats-io/nats.java) - (The Apache License, Version 2.0) jnats-json (io.nats:jnats-json:2.0.1 - https://github.com/nats-io/nats.java.json) - (The Apache License, Version 2.0) jwt-java (io.nats:jwt-java:2.2.0 - https://github.com/nats-io/jwt.java) - (The Apache License, Version 2.0) nkeys-java (io.nats:nkeys-java:2.1.1 - https://github.com/nats-io/nkeys.java) (Apache License, Version 2.0) Netty/Buffer (io.netty:netty-buffer:4.2.15.Final - https://netty.io/) (Apache License, Version 2.0) Netty/Codec (io.netty:netty-codec:4.2.15.Final - https://netty.io/) (Apache License, Version 2.0) Netty/Codec/Base (io.netty:netty-codec-base:4.2.15.Final - https://netty.io/) diff --git a/src/control-plane-services/instance-cluster-management/icms-core/BUILD.bazel b/src/control-plane-services/instance-cluster-management/icms-core/BUILD.bazel index 357743cff..3c3ad613b 100644 --- a/src/control-plane-services/instance-cluster-management/icms-core/BUILD.bazel +++ b/src/control-plane-services/instance-cluster-management/icms-core/BUILD.bazel @@ -45,7 +45,6 @@ ICMS_CORE_DEPS = [ "@nv_third_party_deps//:io_micrometer_micrometer_registry_prometheus", "@nv_third_party_deps//:io_micrometer_micrometer_tracing", "@nv_third_party_deps//:io_nats_jnats", - "@nv_third_party_deps//:io_nats_jwt_java", "@nv_third_party_deps//:io_netty_netty_transport", "@nv_third_party_deps//:io_netty_netty_handler", "@nv_third_party_deps//:io_opentelemetry_opentelemetry_exporter_otlp", diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/configuration/nats/NatsConfiguration.java b/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/configuration/nats/NatsConfiguration.java new file mode 100644 index 000000000..954429104 --- /dev/null +++ b/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/configuration/nats/NatsConfiguration.java @@ -0,0 +1,369 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * 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 + * + * http://www.apache.org/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. + */ +package com.nvidia.icms.configuration.nats; + +import static org.springframework.beans.factory.config.BeanDefinition.SCOPE_PROTOTYPE; + +import io.micrometer.core.instrument.MeterRegistry; +import io.micrometer.tracing.Tracer; +import io.nats.client.Connection; +import io.nats.client.Connection.Status; +import io.nats.client.ConnectionListener; +import io.nats.client.ConnectionListener.Events; +import io.nats.client.ErrorListener; +import io.nats.client.ForceReconnectOptions; +import io.nats.client.JetStream; +import io.nats.client.JetStreamManagement; +import io.nats.client.Nats; +import io.nats.client.Options; +import io.nats.client.Statistics; +import io.nats.client.impl.TracedNatsConnection; +import java.io.IOException; +import java.time.Duration; +import java.util.AbstractCollection; +import java.util.Arrays; +import java.util.Collection; +import java.util.Iterator; +import java.util.List; +import java.util.Objects; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.ToDoubleFunction; +import java.util.function.ToLongFunction; +import lombok.extern.slf4j.Slf4j; +import org.apache.commons.lang3.RandomUtils; +import org.springframework.beans.factory.InitializingBean; +import org.springframework.context.ApplicationContext; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; +import org.springframework.context.annotation.Scope; +import org.springframework.stereotype.Service; + +@Slf4j +@Configuration(proxyBeanMethods = false) +public final class NatsConfiguration { + + private static Connection connect(Options options, Tracer tracer) + throws IOException, InterruptedException { + TracedNatsConnection conn = new TracedNatsConnection(options, tracer); + conn.connect(false); + return conn; + } + + @Bean + @Scope(SCOPE_PROTOTYPE) + public Connection natsConnection( + NatsConfigurationProperties natsProperties, + // only required if the auth callout is enabled to establish ordering + Tracer tracer) + throws IOException, InterruptedException { + + + var options = createDefaultOptions(natsProperties); + return connect(options, tracer); + } + + /** + * Creates default connection options for the NATS client. + * Allows customization of reconnection behavior and other connection parameters. + * + * @return Options object containing the connection configuration. + */ + Options createDefaultOptions(NatsConfigurationProperties natsConfigurationProperties) { + Options.Builder builder = new Options.Builder() + // Set the NATS server URL + .server(natsConfigurationProperties.getNatsUrl()) + // Set connection timeout + .connectionTimeout(natsConfigurationProperties.getConnectionTimeout()) + // Set ping interval + .pingInterval(natsConfigurationProperties.getPingInterval()) + .useDispatcherWithExecutor() + .reconnectWait(natsConfigurationProperties.getReconnectWait()) + .errorListener(new LoggingNatsErrorListener()) + .connectionListener(getNatsConnectionListener()); + + if (natsConfigurationProperties.getReconnectJitter().isPositive()) { + Duration reconnectJitter = natsConfigurationProperties.getReconnectJitter(); + builder.reconnectJitter(reconnectJitter).reconnectJitterTls(reconnectJitter); + } + + // Configure reconnection behavior based on the allowReconnect flag + if (!natsConfigurationProperties.isReconnectAllowed()) { + builder = builder.noReconnect(); // Disable reconnections + } else { + builder = builder.maxReconnects(-1); // Allow unlimited reconnections + } + if (natsConfigurationProperties.getNkeySeed().isPresent()) { + var authHandler = Nats.staticCredentials(null, + natsConfigurationProperties.getNkeySeed() + .get().toCharArray()); + builder = builder.authHandler(authHandler); + } + + return builder.build(); + } + + /** + * Provides a connection listener to handle NATS connection events. + * Specifically handles "LAME_DUCK" events by forcing a reconnection with a jittered delay. + * + * @return ConnectionListener instance to handle connection events. + */ + private ConnectionListener getNatsConnectionListener() { + return (conn, type) -> { + log.info("nats connection event {} {}", type, + conn.getServerInfo()); + if (type == Events.LAME_DUCK) { + CompletableFuture.runAsync(() -> { + try { + // jitter + Thread.sleep(RandomUtils.secure().randomInt(0, 5000)); + // this may cause issues, but hopefully the active force + // reconnection is a smaller error window than waiting to get + // booted and detecting it normally. + log.info("client id {} force reconnecting to nats", + conn.getServerInfo().getClientId()); + conn.forceReconnect(ForceReconnectOptions.builder() + .flush(Duration.ofSeconds(5)) + .build()); + log.info("client id {} reconnected to nats", + conn.getServerInfo().getClientId()); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + log.warn("client id {} reconnect interrupted", + conn.getServerInfo().getClientId(), e); + throw new RuntimeException(e); + } catch (Exception e) { + log.warn("client id {} failed to reconnect to nats", + conn.getServerInfo().getClientId(), e); + throw new RuntimeException(e); + } + }); + } + }; + } + + /** + * Custom error listener for logging NATS errors and exceptions. + */ + @Slf4j + static class LoggingNatsErrorListener implements ErrorListener { + + /** + * Logs errors that occur during the NATS connection lifecycle. + * + * @param conn The NATS connection where the error occurred. + * @param error The error message. + */ + @Override + public void errorOccurred(Connection conn, String error) { + log.error("NATS error occurred {} {}", conn.getServerInfo(), error); + } + + /** + * Logs exceptions that occur during the NATS connection lifecycle. + * + * @param conn The NATS connection where the exception occurred. + * @param exp The exception instance. + */ + @Override + public void exceptionOccurred(Connection conn, Exception exp) { + log.error("For NATS or server: {} error {} : ", conn.getServerInfo(), exp.getMessage(), exp); + } + } + + @Slf4j + @Service + public static class FixedNatsPool implements AutoCloseable { + + private final Connection[] connections; + private final JetStream[] jetStreams; + private final JetStreamManagement[] jetStreamManagements; + private final ApplicationContext applicationContext; + private final AtomicInteger index = new AtomicInteger(); + + public FixedNatsPool(ApplicationContext applicationContext, + NatsConfigurationProperties natsProperties) { + int poolSize = natsProperties.isNatsEnabled() + ? Math.min(Runtime.getRuntime().availableProcessors(), + natsProperties.getMaxPoolSize()) + : 0; + this.applicationContext = applicationContext; + this.connections = new Connection[poolSize]; + this.jetStreams = new JetStream[poolSize]; + this.jetStreamManagements = new JetStreamManagement[poolSize]; + } + + @Override + public void close() throws Exception { + for (Connection connection : connections) { + if (connection != null) { + connection.close(); + } + } + } + + private int nextIndex() { + if (connections.length == 0) { + throw new IllegalStateException("NATS is not enabled"); + } + return Math.floorMod(index.getAndIncrement(), connections.length); + } + + public Connection borrowConnection() throws IOException, InterruptedException { + return connection(nextIndex()); + } + + public JetStream borrowJetStream() throws IOException, InterruptedException { + int slot = nextIndex(); + connection(slot); + return jetStreams[slot]; + } + + public JetStreamManagement borrowJetStreamManagement() + throws IOException, InterruptedException { + int slot = nextIndex(); + connection(slot); + return jetStreamManagements[slot]; + } + + private synchronized Connection connection(int slot) + throws IOException, InterruptedException { + Connection connection = connections[slot]; + if (connection == null || connection.getStatus() == Status.CLOSED) { + connection = applicationContext.getBean(Connection.class); + try { + connection.RTT(); + JetStream jetStream = connection.jetStream(); + JetStreamManagement jetStreamManagement = + connection.jetStreamManagement(); + connections[slot] = connection; + jetStreams[slot] = jetStream; + jetStreamManagements[slot] = jetStreamManagement; + } catch (IOException | RuntimeException e) { + try { + connection.close(); + } catch (InterruptedException closeError) { + Thread.currentThread().interrupt(); + e.addSuppressed(closeError); + } + throw e; + } + } + return connection; + } + + public boolean healthy() { + if (connections.length == 0) { + return false; + } + if (Arrays.stream(connections).allMatch(Objects::isNull)) { + try { + connection(0); + } catch (IOException | InterruptedException | RuntimeException e) { + if (e instanceof InterruptedException) { + Thread.currentThread().interrupt(); + } + log.warn("Unable to initialize NATS connection during health check", e); + return false; + } + } + boolean initialized = false; + for (Connection connection : connections) { + if (connection == null) { + continue; + } + initialized = true; + if (connection.getStatus() != Status.CONNECTED) { + log.warn("Unhealthy NATS connection {}", connection.getServerInfo()); + return false; + } + } + return initialized; + } + + Collection statistics() { + return new AbstractCollection<>() { + @Override + public Iterator iterator() { + return Arrays.stream(connections) + .filter(Objects::nonNull) + .map(Connection::getStatistics) + .iterator(); + } + + @Override + public int size() { + return (int) Arrays.stream(connections).filter(Objects::nonNull).count(); + } + }; + } + } + + @Configuration(proxyBeanMethods = false) + static class NatsMetricsConfiguration implements InitializingBean { + + private final Collection statistics; + private final MeterRegistry meterRegistry; + + NatsMetricsConfiguration(FixedNatsPool fixedNatsPool, MeterRegistry meterRegistry) { + this.meterRegistry = meterRegistry; + this.statistics = fixedNatsPool.statistics(); + } + + @Override + public void afterPropertiesSet() { + meterRegistry.more().counter("nats.pings", List.of(), + statistics, sumProperties(Statistics::getPings)); + meterRegistry.more().counter("nats.reconnects", List.of(), + statistics, sumProperties(Statistics::getReconnects)); + meterRegistry.more().counter("nats.dropped", List.of(), + statistics, sumProperties(Statistics::getDroppedCount)); + meterRegistry.more().counter("nats.oks", List.of(), + statistics, sumProperties(Statistics::getOKs)); + meterRegistry.more().counter("nats.errs", List.of(), + statistics, sumProperties(Statistics::getErrs)); + meterRegistry.more().counter("nats.exceptions", List.of(), + statistics, sumProperties(Statistics::getExceptions)); + meterRegistry.more().counter("nats.requests.sent", List.of(), + statistics, sumProperties(Statistics::getRequestsSent)); + meterRegistry.more().counter("nats.replies.received", List.of(), statistics, + sumProperties(Statistics::getRepliesReceived)); + meterRegistry.more().counter("nats.replies.received.duplicate", List.of(), statistics, + sumProperties(Statistics::getDuplicateRepliesReceived)); + meterRegistry.more().counter("nats.replies.received.orphan", List.of(), statistics, + sumProperties(Statistics::getOrphanRepliesReceived)); + meterRegistry.more().counter("nats.msgs.in", List.of(), + statistics, sumProperties(Statistics::getInMsgs)); + meterRegistry.more().counter("nats.msgs.out", List.of(), + statistics, sumProperties(Statistics::getOutMsgs)); + meterRegistry.more().counter("nats.bytes.in", List.of(), + statistics, sumProperties(Statistics::getInBytes)); + meterRegistry.more().counter("nats.bytes.out", List.of(), + statistics, sumProperties(Statistics::getOutBytes)); + meterRegistry.more().counter("nats.flush", List.of(), + statistics, sumProperties(Statistics::getFlushCounter)); + meterRegistry.gauge("nats.requests.outstanding", statistics, + sumProperties(Statistics::getOutstandingRequests)); + } + + private static ToDoubleFunction> sumProperties( + ToLongFunction propertyExtractor) { + return collection -> collection.stream().mapToLong(propertyExtractor).sum(); + } + } +} diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/configuration/bean/NatsConfigurationProperties.java b/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/configuration/nats/NatsConfigurationProperties.java similarity index 95% rename from src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/configuration/bean/NatsConfigurationProperties.java rename to src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/configuration/nats/NatsConfigurationProperties.java index 261d93dc5..86f405fa8 100644 --- a/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/configuration/bean/NatsConfigurationProperties.java +++ b/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/configuration/nats/NatsConfigurationProperties.java @@ -14,7 +14,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package com.nvidia.icms.configuration.bean; +package com.nvidia.icms.configuration.nats; import java.time.Duration; import java.util.Optional; @@ -33,8 +33,8 @@ public class NatsConfigurationProperties { private boolean natsEnabled; private String natsUrl; + private int maxPoolSize = 8; private boolean createNatsStreams; - private boolean createNatsConsumers; private Duration connectionTimeout = Duration.ZERO; private Duration pingInterval = Duration.ZERO; private Duration reconnectWait = Duration.ZERO; diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/configuration/nats/NatsHealthConfiguration.java b/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/configuration/nats/NatsHealthConfiguration.java new file mode 100644 index 000000000..32697dec1 --- /dev/null +++ b/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/configuration/nats/NatsHealthConfiguration.java @@ -0,0 +1,42 @@ +package com.nvidia.icms.configuration.nats; + +import com.nvidia.icms.configuration.nats.NatsConfiguration.FixedNatsPool; +import java.util.Map; +import lombok.RequiredArgsConstructor; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.boot.health.autoconfigure.contributor.CompositeHealthContributorConfiguration; +import org.springframework.boot.health.contributor.AbstractHealthIndicator; +import org.springframework.boot.health.contributor.Health; +import org.springframework.boot.health.contributor.HealthContributor; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; + +@Configuration(proxyBeanMethods = false) +@ConditionalOnProperty(prefix = "icms.nats", name = "nats-enabled", havingValue = "true") +public class NatsHealthConfiguration extends + CompositeHealthContributorConfiguration { + + NatsHealthConfiguration() { + super(NatsHealthIndicator::new); + } + + @RequiredArgsConstructor + public static class NatsHealthIndicator extends AbstractHealthIndicator { + + private final FixedNatsPool connection; + + @Override + protected void doHealthCheck(Health.Builder builder) { + if (connection.healthy()) { + builder.up(); + } else { + builder.down(); + } + } + } + + @Bean + HealthContributor natsHealthContributor(Map connections) { + return createContributor(connections); + } +} diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/CoreNatsStreamRegistrar.java b/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/CoreNatsStreamRegistrar.java deleted file mode 100644 index fd1df2ca9..000000000 --- a/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/CoreNatsStreamRegistrar.java +++ /dev/null @@ -1,54 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - * - * 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 - * - * http://www.apache.org/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. - */ -package com.nvidia.icms.outbound.nats; - -import java.util.List; -import org.springframework.stereotype.Component; - -/** - * NVCA JetStream streams that ICMS-core always needs. Lives in core because the NVCA - * pipeline is part of the always-on, non BYOC deployment surface. - * - *

Stream + subject string values are byte-identical to the constants previously declared inline - * in {@code NatsStreamManager} so no NATS server-side state changes are required by R8.

- */ -@Component -public class CoreNatsStreamRegistrar implements NatsStreamRegistrar { - - public static final String CREATE_NVCA_STREAM_NAME = "CreateNvcaFunctionTaskStream"; - public static final String TERMINATE_NVCA_STREAM_NAME = "TerminateNvcaStream"; - - private static final String CREATE_NVCA_STREAM_SUBJECT = "Create.NVCA.>"; - private static final String CREATE_NVCA_CONSUMER_SUBJECT = "Create.NVCA.*.*.*.*"; - private static final String TERMINATE_NVCA_STREAM_SUBJECT = "Terminate.NVCA.>"; - private static final String TERMINATE_NVCA_CONSUMER_SUBJECT = "Terminate.NVCA.*"; - - private static final List DEFINITIONS = List.of( - new NatsStreamDefinition( - CREATE_NVCA_STREAM_NAME, - CREATE_NVCA_STREAM_SUBJECT, - CREATE_NVCA_CONSUMER_SUBJECT), - new NatsStreamDefinition( - TERMINATE_NVCA_STREAM_NAME, - TERMINATE_NVCA_STREAM_SUBJECT, - TERMINATE_NVCA_CONSUMER_SUBJECT)); - - @Override - public List getStreamDefinitions() { - return DEFINITIONS; - } -} diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/NatsConnectionFactory.java b/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/NatsConnectionFactory.java deleted file mode 100644 index a341be583..000000000 --- a/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/NatsConnectionFactory.java +++ /dev/null @@ -1,239 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - * - * 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 - * - * http://www.apache.org/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. - */ -package com.nvidia.icms.outbound.nats; - -import com.nvidia.icms.configuration.bean.NatsConfigurationProperties; -import io.nats.client.Connection; -import io.nats.client.ConnectionListener; -import io.nats.client.ConnectionListener.Events; -import io.nats.client.ErrorListener; -import io.nats.client.ForceReconnectOptions; -import io.nats.client.Nats; -import io.nats.client.Options; -import java.io.IOException; -import java.time.Duration; -import java.util.concurrent.CompletableFuture; -import lombok.extern.slf4j.Slf4j; -import org.apache.commons.lang3.RandomUtils; -import org.springframework.stereotype.Component; - -/** - * Factory class for managing connections to a NATS server. - * This class provides methods to establish, configure, and manage a connection to the NATS server. - * It ensures a single connection instance is reused and handles reconnection logic when necessary. - */ -@Component -@Slf4j -public class NatsConnectionFactory { - - private final NatsConfigurationProperties natsConfigurationProperties; - - // Instance of the NATS connection - // volatile needs to safely read without synchronization in getCachedConnection method - private volatile Connection natsConnection; - - /** - * Constructor to initialize the NATS connection factory with configuration properties. - * - * @param natsConfigurationProperties Configuration properties for NATS connection. - */ - public NatsConnectionFactory(NatsConfigurationProperties natsConfigurationProperties) { - this.natsConfigurationProperties = natsConfigurationProperties; - } - - /** - * Establishes a connection to the NATS server. - * If a connection already exists, it reuses the existing connection. - * - * @return Connection object representing the active connection to the NATS server. - * @throws IOException If an I/O error occurs during the connection process. - * @throws InterruptedException If the connection attempt is interrupted. - */ - public synchronized Connection createConnectionIfNeeded() throws IOException, InterruptedException { - if (natsConnection == null || natsConnection.getStatus() == Connection.Status.CLOSED) { - natsConnection = connectToNats(); - } - return natsConnection; - } - - /** - * Drops the cached connection so the next {@link #createConnectionIfNeeded()} call - * re-handshakes. Used during startup retry when a prior connect failed - * (e.g. NATS auth-callout was not yet ready) — the cached failure - * connection is unusable and must be discarded before the retry. - */ - public synchronized void resetConnection() { - if (natsConnection != null) { - try { - natsConnection.close(); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - log.warn("Interrupted while closing stale NATS connection", e); - } finally { - // finally block will be executed even if we interrupt the tread - natsConnection = null; - } - } - } - - /** - * Removes a terminal connection from the cache without disturbing a newer connection that - * may already have replaced it. - */ - synchronized void invalidateClosedConnection(Connection connection) { - if (natsConnection == connection && connection.getStatus() == Connection.Status.CLOSED) { - natsConnection = null; - } - } - - /** - * Returns the current connection without creating one. Intended for health reporting. - */ - Connection getCachedConnection() { - return natsConnection; - } - - Connection connectToNats() - throws IOException, InterruptedException { - if (natsConfigurationProperties.isNatsEnabled()) { - Options options = createDefaultOptions(natsConfigurationProperties.isReconnectAllowed()); - - try { - log.info("NATS: Connecting to server {}", options.getServers().getFirst()); - Connection connection = Nats.connect(options); - log.info("NATS: Successfully connected to server {}", - options.getServers().getFirst()); - return connection; - } catch (IOException e) { - log.error("NATS: Failed to connect to server {}: {}", options.getServers().getFirst(), e.getMessage(), e); - throw e; - } - } - else { - return null; - } - } - - /** - * Creates default connection options for the NATS client. - * Allows customization of reconnection behavior and other connection parameters. - * - * @param allowReconnect Boolean flag to enable or disable reconnections. - * @return Options object containing the connection configuration. - */ - Options createDefaultOptions(Boolean allowReconnect) { - Options.Builder builder = new Options.Builder() - // Set the NATS server URL - .server(natsConfigurationProperties.getNatsUrl()) - // Set connection timeout - .connectionTimeout(natsConfigurationProperties.getConnectionTimeout()) - // Set ping interval - .pingInterval(natsConfigurationProperties.getPingInterval()) - .useDispatcherWithExecutor() - .reconnectWait(natsConfigurationProperties.getReconnectWait()) - .errorListener(new LoggingNatsErrorListener()) - .connectionListener(getNatsConnectionListener()); - - if (natsConfigurationProperties.getReconnectJitter().isPositive()) { - Duration reconnectJitter = natsConfigurationProperties.getReconnectJitter(); - builder.reconnectJitter(reconnectJitter).reconnectJitterTls(reconnectJitter); - } - - // Configure reconnection behavior based on the allowReconnect flag - if (!allowReconnect) { - builder = builder.noReconnect(); // Disable reconnections - } else { - builder = builder.maxReconnects(-1); // Allow unlimited reconnections - } - if (natsConfigurationProperties.getNkeySeed().isPresent()) { - var authHandler = Nats.staticCredentials(null, - natsConfigurationProperties.getNkeySeed() - .get().toCharArray()); - builder = builder.authHandler(authHandler); - } - - return builder.build(); - } - - /** - * Provides a connection listener to handle NATS connection events. - * Specifically handles "LAME_DUCK" events by forcing a reconnection with a jittered delay. - * - * @return ConnectionListener instance to handle connection events. - */ - private ConnectionListener getNatsConnectionListener() { - return (connection, events) -> { - log.info("NATS connection event {} {}", connection.getServerInfo(), events); - if (events == Events.CLOSED) { - invalidateClosedConnection(connection); - } else if (events == Events.LAME_DUCK) { - CompletableFuture.runAsync(() -> { - try { - // Add jitter to avoid simultaneous reconnections - Thread.sleep(RandomUtils.nextInt(0, 5000)); - log.info("Client ID {} force reconnecting to NATS", - connection.getServerInfo().getClientId()); - connection.forceReconnect( - ForceReconnectOptions.builder() - .flush(natsConfigurationProperties.getForceReconnectFlush()) - .build()); - log.info("Client ID {} reconnected to NATS", - connection.getServerInfo().getClientId()); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - log.warn("Client ID {} reconnect interrupted", - connection.getServerInfo().getClientId(), e); - throw new RuntimeException(e); - } catch (Exception e) { - log.warn("Client ID {} failed to reconnect to NATS", - connection.getServerInfo().getClientId(), e); - throw new RuntimeException(e); - } - }); - } - }; - } - - /** - * Custom error listener for logging NATS errors and exceptions. - */ - @Slf4j - static class LoggingNatsErrorListener implements ErrorListener { - - /** - * Logs errors that occur during the NATS connection lifecycle. - * - * @param conn The NATS connection where the error occurred. - * @param error The error message. - */ - @Override - public void errorOccurred(Connection conn, String error) { - log.error("NATS error occurred {} {}", conn.getServerInfo(), error); - } - - /** - * Logs exceptions that occur during the NATS connection lifecycle. - * - * @param conn The NATS connection where the exception occurred. - * @param exp The exception instance. - */ - @Override - public void exceptionOccurred(Connection conn, Exception exp) { - log.error("For NATS or server: {} error {} : ", conn.getServerInfo(), exp.getMessage(), exp); - } - } -} diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/NatsHealthIndicator.java b/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/NatsHealthIndicator.java deleted file mode 100644 index c23daa649..000000000 --- a/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/NatsHealthIndicator.java +++ /dev/null @@ -1,58 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - * - * 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 - * - * http://www.apache.org/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. - */ -package com.nvidia.icms.outbound.nats; - -import io.nats.client.Connection; -import org.apache.commons.lang3.StringUtils; -import org.springframework.boot.health.contributor.Health; -import org.springframework.boot.health.contributor.HealthIndicator; -import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; -import org.springframework.stereotype.Component; - -/** Reports whether the NATS connection used for NVCA messaging is currently usable. */ -@Component -@ConditionalOnProperty(prefix = "icms.nats", name = "nats-enabled", havingValue = "true") -public class NatsHealthIndicator implements HealthIndicator { - - private final NatsConnectionFactory natsConnectionFactory; - - public NatsHealthIndicator(NatsConnectionFactory natsConnectionFactory) { - this.natsConnectionFactory = natsConnectionFactory; - } - - @Override - public Health health() { - Connection connection = natsConnectionFactory.getCachedConnection(); - if (connection == null) { - return Health.down().withDetail("status", "NOT_INITIALIZED").build(); - } - - Connection.Status status = connection.getStatus(); - Health.Builder health = status == Connection.Status.CONNECTED ? Health.up() : Health.down(); - health.withDetail("status", status == null ? "UNKNOWN" : status.name()); - - String connectedUrl = connection.getConnectedUrl(); - if (StringUtils.isNotBlank(connectedUrl)) { - health.withDetail("server", connectedUrl); - } - String lastError = connection.getLastError(); - if (StringUtils.isNotBlank(lastError)) { - health.withDetail("lastError", lastError); - } - return health.build(); - } -} diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/NatsMessageSenderClient.java b/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/NatsMessageSenderClient.java index d6bc92731..74e8a2b25 100644 --- a/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/NatsMessageSenderClient.java +++ b/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/NatsMessageSenderClient.java @@ -16,12 +16,12 @@ */ package com.nvidia.icms.outbound.nats; -import com.nvidia.icms.configuration.bean.NatsConfigurationProperties; +import com.nvidia.icms.configuration.nats.NatsConfiguration.FixedNatsPool; +import com.nvidia.icms.configuration.nats.NatsConfigurationProperties; import com.nvidia.icms.errors.IcmsInternalServerException; import com.nvidia.icms.outbound.sqs.model.byoc.ByocSqsMessageModel; -import com.nvidia.icms.util.GsonCompatMapper; import com.nvidia.icms.outbound.sqs.model.byoc.ByocTerminatePodMessageModel; -import io.nats.client.Connection; +import com.nvidia.icms.util.GsonCompatMapper; import io.nats.client.JetStream; import io.nats.client.impl.NatsMessage; import jakarta.annotation.Nullable; @@ -57,7 +57,7 @@ public enum SendNatsMessageResult { NOT_ENABLED // NATS is not enabled } - private final NatsConnectionFactory natsConnectionFactory; + private final FixedNatsPool fixedNatsPool; private final NatsConfigurationProperties natsConfigurationProperties; /** @@ -235,8 +235,7 @@ private SendNatsMessageResult sendTerminationMessage( SendNatsMessageResult publishMessage( @NotNull String messageBody, @NotNull String subject) { try { - Connection connection = natsConnectionFactory.createConnectionIfNeeded(); - JetStream js = connection.jetStream(); + JetStream js = fixedNatsPool.borrowJetStream(); js.publish(NatsMessage.builder() .subject(subject) .data(messageBody, StandardCharsets.UTF_8) diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/NatsResourceService.java b/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/NatsResourceService.java new file mode 100644 index 000000000..12b9eabff --- /dev/null +++ b/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/NatsResourceService.java @@ -0,0 +1,76 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * 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 + * + * http://www.apache.org/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. + */ +package com.nvidia.icms.outbound.nats; + +import com.nvidia.icms.configuration.nats.NatsConfiguration.FixedNatsPool; +import io.micrometer.core.annotation.Timed; +import io.nats.client.JetStreamApiException; +import io.nats.client.JetStreamManagement; +import io.nats.client.api.StreamConfiguration; +import java.io.IOException; +import lombok.RequiredArgsConstructor; +import lombok.extern.slf4j.Slf4j; +import org.springframework.stereotype.Service; + +@Slf4j +@Service +@RequiredArgsConstructor +public class NatsResourceService { + + private final FixedNatsPool fixedNatsPool; + + @Timed(value = "icms.nats.create.stream") + public void createStream(StreamConfiguration streamConfig) + throws IOException, JetStreamApiException, InterruptedException { + JetStreamManagement jetStreamManagement = fixedNatsPool.borrowJetStreamManagement(); + try { + var streamInfo = jetStreamManagement.getStreamInfo(streamConfig.getName()); + validateStreamConfiguration(streamConfig, streamInfo.getConfiguration()); + return; + } catch (JetStreamApiException e) { + // non-404 related error gets passed back up + if (e.getErrorCode() != 404) { + throw e; + } + // if stream doesn't exist, keep going and try to create + } + try { + // if the stream was created by another server during this gap + // but the config is the same, this call will succeed + jetStreamManagement.addStream(streamConfig); + } catch (JetStreamApiException e) { + if (e.getApiErrorCode() == 10058) { + var streamInfo = jetStreamManagement.getStreamInfo(streamConfig.getName()); + validateStreamConfiguration(streamConfig, streamInfo.getConfiguration()); + return; + } + throw e; + } + } + + private static void validateStreamConfiguration( + StreamConfiguration expected, StreamConfiguration actual) { + if (!expected.getSubjects().equals(actual.getSubjects()) + || expected.getStorageType() != actual.getStorageType() + || expected.getRetentionPolicy() != actual.getRetentionPolicy() + || expected.getMaxMsgs() != actual.getMaxMsgs() + || !expected.getMaxAge().equals(actual.getMaxAge())) { + throw new IllegalStateException( + "NATS stream " + expected.getName() + " has an incompatible configuration"); + } + } +} diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/NatsStreamDefinition.java b/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/NatsStreamDefinition.java deleted file mode 100644 index 604466bca..000000000 --- a/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/NatsStreamDefinition.java +++ /dev/null @@ -1,48 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - * - * 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 - * - * http://www.apache.org/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. - */ -package com.nvidia.icms.outbound.nats; - -import jakarta.validation.constraints.NotBlank; - -/** - * Declarative description of a single JetStream stream and its companion consumer. - * - *

Used by {@link NatsStreamRegistrar} so that modules (core, non-BYOC, …) can declare the streams - * they own without {@link NatsStreamManager} hard-coding any module-specific names. The manager - * iterates {@link NatsStreamRegistrar} beans and creates / validates one stream + one consumer per - * definition.

- * - *

Consumer name convention is {@code streamName + "Consumer"} (preserved verbatim from the - * original {@code NatsStreamManager.createNatsConsumers()} implementation so existing NATS server - * state continues to match).

- * - * @param streamName the JetStream stream name (e.g. {@code CreateNvcaFunctionTaskStream}). - * @param streamSubject the broad subject pattern attached to the stream - * (e.g. {@code Create.NVCA.>}). - * @param consumerSubject the filter subject the durable consumer subscribes to - * (e.g. {@code Create.NVCA.*.*.*.*}). - */ -public record NatsStreamDefinition( - @NotBlank String streamName, - @NotBlank String streamSubject, - @NotBlank String consumerSubject) { - - /** Convention used by {@code NatsStreamManager} since inception. */ - public String consumerName() { - return streamName + "Consumer"; - } -} diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/NatsStreamManager.java b/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/NatsStreamManager.java index 1856b27ed..78fe7ad46 100644 --- a/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/NatsStreamManager.java +++ b/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/NatsStreamManager.java @@ -16,374 +16,106 @@ */ package com.nvidia.icms.outbound.nats; -import com.nvidia.icms.configuration.bean.NatsConfigurationProperties; -import io.nats.client.ConsumerContext; -import io.nats.client.JetStreamApiException; -import io.nats.client.JetStreamManagement; -import io.nats.client.StreamContext; -import io.nats.client.api.AckPolicy; -import io.nats.client.api.ConsumerConfiguration; +import com.nvidia.icms.configuration.nats.NatsConfigurationProperties; import io.nats.client.api.RetentionPolicy; import io.nats.client.api.StorageType; import io.nats.client.api.StreamConfiguration; -import io.nats.client.api.StreamInfo; import jakarta.annotation.PostConstruct; -import java.io.IOException; import java.time.Duration; -import java.util.ArrayList; import java.util.List; +import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; import org.springframework.stereotype.Component; -/** - * The NatsStreamManager class is responsible for managing JetStream streams and consumers on a NATS server. - * - *

This class provides methods to: - *

    - *
  • Validate the existence of required streams and create them if necessary.
    • - *
  • Create new JetStream streams with specific configurations.
    • - *
  • Retrieve information about existing streams.
    • - *
  • Delete streams and manage JetStream configurations using the NATS connection.
    • - *
  • Create and manage consumers for specific streams.
    • - *
- * - *

It integrates with the NATS server using {@link NatsConnectionFactory} and relies on - * {@link NatsConfigurationProperties} for configuration values. - */ @Component +@RequiredArgsConstructor @Slf4j public class NatsStreamManager { - private final NatsConnectionFactory natsConnectionFactory; - private final NatsConfigurationProperties natsConfigurationProperties; - /** - * Every {@link NatsStreamRegistrar} discovered in the application context. icms-core always - * contributes {@code CoreNatsStreamRegistrar} (the NVCA stream pair); a deployment module may - * contribute an additional registrar (e.g. the non-BYOC stream pair). Deployments running only - * icms-core register just the NVCA pair. - */ - private final List streamRegistrars; + public static final String CREATE_NVCA_STREAM_NAME = "CreateNvcaFunctionTaskStream"; + public static final String TERMINATE_NVCA_STREAM_NAME = "TerminateNvcaStream"; - /** - * Constructor for NatsStreamManager. - * - * @param natsConnectionFactory Factory for creating NATS connections. - * @param natsConfigurationProperties Configuration properties for NATS. - * @param streamRegistrars All {@link NatsStreamRegistrar} beans from the context (may be empty - * if no module declares any streams; the constructor accepts an empty - * list in that case). - */ - public NatsStreamManager( - NatsConnectionFactory natsConnectionFactory, - NatsConfigurationProperties natsConfigurationProperties, - List streamRegistrars) { + private static final String CREATE_NVCA_STREAM_SUBJECT = "Create.NVCA.>"; + private static final String TERMINATE_NVCA_STREAM_SUBJECT = "Terminate.NVCA.>"; + private static final int MAX_MESSAGES = 1_000_000; + private static final int MAX_INIT_ATTEMPTS = 60; + private static final Duration INIT_RETRY_DELAY = Duration.ofSeconds(5); - this.natsConnectionFactory = natsConnectionFactory; - this.natsConfigurationProperties = natsConfigurationProperties; - this.streamRegistrars = streamRegistrars == null ? List.of() : streamRegistrars; - } - - /** - * Returns every {@link NatsStreamDefinition} contributed by the registered modules, in - * registrar order. Order is irrelevant to NATS itself but is preserved so logs remain - * deterministic across runs with the same registrar set. - */ - private List getAllStreamDefinitions() { - List merged = new ArrayList<>(); - for (NatsStreamRegistrar registrar : streamRegistrars) { - List registrarDefinitions = registrar.getStreamDefinitions(); - if (registrarDefinitions != null) { - merged.addAll(registrarDefinitions); - } - } - return merged; - } + private final NatsResourceService natsResourceService; + private final NatsConfigurationProperties natsConfigurationProperties; - /** - * Streams + consumers are required for SIS to enqueue function-deploy work - * to the NVCA agent. NATS is brought up in parallel with SIS, and on - * self-managed clusters the auth-callout webhook also has to be ready - * before the NKey-mapped SIS user can connect — so this init is allowed to - * race the NATS connection on startup. Retry with backoff and ultimately - * throw if NATS never becomes usable, so the Spring container fails and - * Kubernetes restarts the pod into a clean window where the streams can - * actually be created. - * - * The runtime re-validation path - * (com.nvidia.icms.scheduled.GlobalNatsStreamValidationTaskController) - * keeps the existing tolerate-and-log behavior — a transient mid-life - * NATS hiccup must not crash SIS. - */ @PostConstruct public void init() { - if (!natsConfigurationProperties.isNatsEnabled()) { - return; - } - boolean wantStreams = natsConfigurationProperties.isCreateNatsStreams(); - boolean wantConsumers = natsConfigurationProperties.isCreateNatsConsumers(); - if (!wantStreams && !wantConsumers) { + if (!natsConfigurationProperties.isNatsEnabled() + || !natsConfigurationProperties.isCreateNatsStreams()) { return; } - final int maxAttempts = 60; - final Duration retryDelay = Duration.ofSeconds(5); Exception lastError = null; - - for (int attempt = 1; attempt <= maxAttempts; attempt++) { + for (int attempt = 1; attempt <= MAX_INIT_ATTEMPTS; attempt++) { try { - if (wantStreams) { - validateNatsStreamsStrict(); - } - if (wantConsumers) { - createNatsConsumersStrict(); - } - log.info( - "NATS streams/consumers initialized on attempt {}/{}", - attempt, - maxAttempts); + validateNatsStreamsStrict(); + log.info("NATS streams initialized on attempt {}/{}", + attempt, MAX_INIT_ATTEMPTS); return; } catch (Exception e) { lastError = e; - log.warn( - "NATS init attempt {}/{} failed: {}; retrying in {}s", - attempt, - maxAttempts, - e.getMessage(), - retryDelay.toSeconds()); - // Drop the cached connection so the next attempt can re-handshake - // with auth-callout once it becomes ready. - natsConnectionFactory.resetConnection(); - if (attempt == maxAttempts) { - break; - } - try { - Thread.sleep(retryDelay.toMillis()); - } catch (InterruptedException ie) { - Thread.currentThread().interrupt(); - throw new IllegalStateException( - "Interrupted during NATS init retry", ie); + log.warn("NATS init attempt {}/{} failed: {}; retrying in {}s", + attempt, MAX_INIT_ATTEMPTS, e.getMessage(), + INIT_RETRY_DELAY.toSeconds()); + if (attempt < MAX_INIT_ATTEMPTS) { + sleepBeforeRetry(); } } } + throw new IllegalStateException( - String.format( - "NATS streams/consumers init failed after %d attempts; last error: %s", - maxAttempts, - lastError == null ? "unknown" : lastError.getMessage()), + String.format("NATS stream init failed after %d attempts; last error: %s", + MAX_INIT_ATTEMPTS, + lastError == null ? "unknown" : lastError.getMessage()), lastError); } - /** - * Strict variant of {@link #validateNatsStreams()} that propagates any - * stream creation/lookup failure so {@link #init()} can retry. The - * lenient variant is retained for the scheduled re-validation path. - */ - private void validateNatsStreamsStrict() throws Exception { - for (NatsStreamDefinition definition : getAllStreamDefinitions()) { - getOrCreateStreamStrict(definition.streamName(), definition.streamSubject()); - } - } - - private StreamInfo getOrCreateStreamStrict(String streamName, String streamSubject) - throws Exception { - StreamInfo info = getStream(streamName); - if (info == null) { - info = createStream(streamName, streamSubject); - } - return info; - } - - /** - * Strict variant of {@link #createNatsConsumers()} that propagates any - * consumer-creation failure so {@link #init()} can retry. - */ - private void createNatsConsumersStrict() throws Exception { - for (NatsStreamDefinition definition : getAllStreamDefinitions()) { - createConsumerStrict( - definition.streamName(), - definition.consumerName(), - definition.consumerSubject()); - } - } - - private ConsumerContext createConsumerStrict( - String streamName, String consumerName, String subject) throws Exception { - StreamContext streamContext = - natsConnectionFactory.createConnectionIfNeeded().getStreamContext(streamName); - ConsumerContext result = streamContext.createOrUpdateConsumer( - ConsumerConfiguration.builder() - .durable(consumerName) - .ackPolicy(AckPolicy.Explicit) - .filterSubject(subject) - .build()); - log.info( - "Consumer {} for stream {} with subject {} was created", - consumerName, - streamName, - subject); - return result; - } - - /** - * Validates the existence of required JetStream streams. - * If a stream does not exist, it will be created with the appropriate configuration. - */ public void validateNatsStreams() { - for (NatsStreamDefinition definition : getAllStreamDefinitions()) { - getOrCreateStream(definition.streamName(), definition.streamSubject()); - } - } - - /** - * Creates consumers for predefined streams with specific configurations. - */ - public void createNatsConsumers() { - for (NatsStreamDefinition definition : getAllStreamDefinitions()) { - createConsumer( - definition.streamName(), - definition.consumerName(), - definition.consumerSubject()); - } - } - - /** - * Creates a consumer for a specific stream and subject. - * - * @param streamName The name of the stream. - * @param consumerName The name of the consumer. - * @param subject The subject filter for the consumer. - * @return The created ConsumerContext, or null if an error occurs. - */ - ConsumerContext createConsumer(String streamName, String consumerName, String subject) { - try { - StreamContext streamContext = natsConnectionFactory.createConnectionIfNeeded() - .getStreamContext(streamName); - - ConsumerContext result = streamContext.createOrUpdateConsumer( - ConsumerConfiguration.builder() - .durable(consumerName) - .ackPolicy(AckPolicy.Explicit) - .filterSubject(subject) - .build()); - log.info("Consumer {} for stream {} with subject {} was created", consumerName, - streamName, subject); - return result; - } catch (Exception e) { - log.error("Error creating consumer {} for stream {} with subject {}: {}", consumerName, - streamName, subject, e.getMessage(), e); - } - return null; - } - - /** - * Creates a new JetStream stream on the NATS server with the specified name and subject. - * - * @param streamName The name of the stream to create. - * @param streamSubject The subject to associate with the stream. - * @return StreamInfo object containing details about the created stream. - * @throws IOException If an I/O error occurs during stream creation. - * @throws JetStreamApiException If there is an error in the JetStream API. - * @throws InterruptedException If the operation is interrupted. - */ - public StreamInfo createStream(String streamName, String streamSubject) - throws IOException, JetStreamApiException, InterruptedException { - try { - JetStreamManagement jsm = getJetStreamManagement(); - - StreamConfiguration streamConfig = StreamConfiguration.builder() - .name(streamName) - .subjects(streamSubject) - .storageType(StorageType.Memory) - .retentionPolicy(RetentionPolicy.WorkQueue) - .maxMessages(1000_000) - .maxAge(natsConfigurationProperties.getMessageTtl()) - .build(); - - StreamInfo streamInfo = jsm.addStream(streamConfig); - log.info("Stream {} with subject {} was created", streamName, streamSubject); - return streamInfo; - - } catch (IOException | JetStreamApiException | InterruptedException e) { - log.error("Error creating stream {} with subject {}: {}", streamName, streamSubject, - e.getMessage(), e); - throw e; + for (var streamConfiguration : streamConfigurations()) { + try { + natsResourceService.createStream(streamConfiguration); + } catch (Exception e) { + log.error("Error creating stream {}: {}", streamConfiguration.getName(), + e.getMessage(), e); + } } } - /** - * Deletes an existing JetStream stream by its name. - * - * @param streamName The name of the stream to delete. - * @throws IOException If an I/O error occurs during the operation. - * @throws JetStreamApiException If there is an error in the JetStream API. - * @throws InterruptedException If the operation is interrupted. - */ - public void deleteStream(String streamName) - throws IOException, JetStreamApiException, InterruptedException { - try { - JetStreamManagement jsm = getJetStreamManagement(); - jsm.deleteStream(streamName); - log.info("Stream {} was deleted", streamName); - } catch (IOException | JetStreamApiException | InterruptedException e) { - log.error("Error deleting stream {}: {}", streamName, e.getMessage(), e); - throw e; + private void validateNatsStreamsStrict() throws Exception { + for (var streamConfiguration : streamConfigurations()) { + natsResourceService.createStream(streamConfiguration); } } - /** - * Retrieves information about an existing stream by its name. - * - * @param streamName The name of the stream to retrieve. - * @return StreamInfo object containing details about the stream, or null if the stream does not exist. - * @throws IOException If an I/O error occurs during the operation. - * @throws InterruptedException If the operation is interrupted. - */ - public StreamInfo getStream(String streamName) - throws IOException, InterruptedException { - try { - JetStreamManagement jsm = getJetStreamManagement(); - return jsm.getStreamInfo(streamName); - } catch (JetStreamApiException e) { - log.warn("Stream {} does not exist: {}", streamName, e.getMessage()); - return null; - } + private List streamConfigurations() { + return List.of( + streamConfiguration(CREATE_NVCA_STREAM_NAME, CREATE_NVCA_STREAM_SUBJECT), + streamConfiguration(TERMINATE_NVCA_STREAM_NAME, TERMINATE_NVCA_STREAM_SUBJECT)); } - /** - * Retrieves an existing stream or creates a new one if it does not exist. - * - * @param streamName The name of the stream to retrieve or create. - * @param streamSubject The subject to associate with the stream if it is created. - * @return StreamInfo object containing details about the retrieved or created stream. - */ - public StreamInfo getOrCreateStream(String streamName, String streamSubject) { - StreamInfo streamInfo = null; - try { - streamInfo = getStream(streamName); - if (streamInfo == null) { - streamInfo = createStream(streamName, streamSubject); - } - } catch (Exception e) { - log.error("Error re-creating stream {} with subject {}: {}", streamName, - streamSubject, e.getMessage(), e); - } - return streamInfo; + private StreamConfiguration streamConfiguration(String name, String subject) { + return StreamConfiguration.builder() + .name(name) + .subjects(subject) + .storageType(StorageType.Memory) + .retentionPolicy(RetentionPolicy.WorkQueue) + .maxMessages(MAX_MESSAGES) + .maxAge(natsConfigurationProperties.getMessageTtl()) + .build(); } - /** - * Retrieves the JetStreamManagement context from the NATS connection. - * - * @return JetStreamManagement object for managing JetStream streams. - * @throws IOException If an I/O error occurs while obtaining the context. - * @throws InterruptedException If the operation is interrupted. - */ - private JetStreamManagement getJetStreamManagement() - throws IOException, InterruptedException { + private void sleepBeforeRetry() { try { - return natsConnectionFactory.createConnectionIfNeeded().jetStreamManagement(); - } catch (IOException | InterruptedException e) { - log.error("Error getting JetStreamManagement: {}", e.getMessage(), e); - throw e; + Thread.sleep(INIT_RETRY_DELAY.toMillis()); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException("Interrupted during NATS init retry", e); } } } diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/NatsStreamRegistrar.java b/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/NatsStreamRegistrar.java deleted file mode 100644 index ce25be9b2..000000000 --- a/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/outbound/nats/NatsStreamRegistrar.java +++ /dev/null @@ -1,40 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - * - * 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 - * - * http://www.apache.org/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. - */ -package com.nvidia.icms.outbound.nats; - -import java.util.List; - -/** - * SPI letting each module contribute the NATS JetStream streams + consumers it owns. - * - *

Spring auto-collects every {@code NatsStreamRegistrar} bean in the application context; the - * {@link NatsStreamManager} iterates them on startup (strict path) and during periodic - * re-validation (lenient path) and creates / re-creates every declared stream and consumer.

- * - *

This indirection keeps provider-specific stream names out of core code. Each module owns its - * own stream names + subjects: the NVCA pair lives in {@link CoreNatsStreamRegistrar} (always - * present), and a deployment module may contribute an additional pair (e.g. the non-BYOC streams, - * only loaded when that module is on the classpath).

- */ -public interface NatsStreamRegistrar { - - /** - * Returns the streams this registrar wants {@link NatsStreamManager} to manage. Implementations - * should return a stable, immutable list; the manager treats the list as read-only. - */ - List getStreamDefinitions(); -} diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/scheduled/GlobalNatsStreamValidationTaskController.java b/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/scheduled/GlobalNatsStreamValidationTaskController.java index 7cdb63ded..ea1257f63 100644 --- a/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/scheduled/GlobalNatsStreamValidationTaskController.java +++ b/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/scheduled/GlobalNatsStreamValidationTaskController.java @@ -27,7 +27,7 @@ import com.google.common.base.Stopwatch; import static com.nvidia.icms.configuration.SchedulingConfiguration.SCHEDULED_JOBS_PROFILES; -import com.nvidia.icms.configuration.bean.NatsConfigurationProperties; +import com.nvidia.icms.configuration.nats.NatsConfigurationProperties; import com.nvidia.icms.outbound.nats.NatsStreamManager; import com.nvidia.icms.service.LockProviderService; import com.nvidia.icms.service.telemetry.TelemetryEventClient; diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/service/InstanceServiceHelper.java b/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/service/InstanceServiceHelper.java index 59701f199..0e086d02d 100644 --- a/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/service/InstanceServiceHelper.java +++ b/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/com/nvidia/icms/service/InstanceServiceHelper.java @@ -18,7 +18,7 @@ import tools.jackson.core.JacksonException; import tools.jackson.databind.ObjectMapper; -import com.nvidia.icms.configuration.bean.NatsConfigurationProperties; +import com.nvidia.icms.configuration.nats.NatsConfigurationProperties; import com.nvidia.icms.configuration.bean.IcmsConfigurationProperties; import com.nvidia.icms.configuration.byoc.ByocConfigurationProperties; import com.nvidia.icms.configuration.nvca.NvcaConfigurationProperties; diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/io/nats/client/impl/TracedNatsConnection.java b/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/io/nats/client/impl/TracedNatsConnection.java new file mode 100644 index 000000000..52761fc55 --- /dev/null +++ b/src/control-plane-services/instance-cluster-management/icms-core/src/main/java/io/nats/client/impl/TracedNatsConnection.java @@ -0,0 +1,82 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * 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 + * + * http://www.apache.org/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. + */ +package io.nats.client.impl; + +import io.micrometer.observation.annotation.Observed; +import io.micrometer.tracing.Tracer; +import io.nats.client.Message; +import io.nats.client.Options; +import io.nats.client.support.NatsRequestCompletableFuture.CancelAction; +import java.io.IOException; +import java.time.Duration; +import java.util.Map; +import java.util.concurrent.CompletableFuture; + +public class TracedNatsConnection extends NatsConnection { + + private final Tracer tracer; + + private static final String OBSERVATION_PUBLISH = "nats-publish-internal"; + private static final String OBSERVATION_REQUEST_FUTURE = "nats-request-future-internal"; + private static final String TAG_SUBJECT = "subject"; + private static final String TAG_REPLY_TO = "replyTo"; + private static final String TAG_FUTURE_TIMEOUT = "futureTimeout"; + + public TracedNatsConnection(Options options, Tracer tracer) { + super(options); + this.tracer = tracer; + } + + public void connect(boolean reconnectOnConnect) + throws InterruptedException, IOException { + super.connect(reconnectOnConnect); + } + + @Observed(name = OBSERVATION_PUBLISH) + @Override + void publishInternal( + String subject, String replyTo, Headers headers, + byte[] data, boolean validateSubRep, boolean flushImmediatelyAfterPublish) { + addTagsToCurrentSpan(tracer, Map.of( + TAG_SUBJECT, subject, + TAG_REPLY_TO, String.valueOf(replyTo) + )); + super.publishInternal(subject, replyTo, headers, data, validateSubRep, + flushImmediatelyAfterPublish); + } + + @Observed(name = OBSERVATION_REQUEST_FUTURE) + @Override + CompletableFuture requestFutureInternal( + String subject, Headers headers, byte[] data, + Duration futureTimeout, CancelAction cancelAction, + boolean validateSubjectAndReplyTo, boolean flushImmediatelyAfterPublish) { + addTagsToCurrentSpan(tracer, Map.of( + TAG_SUBJECT, subject, + TAG_FUTURE_TIMEOUT, String.valueOf(futureTimeout) + )); + return super.requestFutureInternal(subject, headers, data, futureTimeout, cancelAction, + validateSubjectAndReplyTo, flushImmediatelyAfterPublish); + } + + private static void addTagsToCurrentSpan(Tracer tracer, Map tags) { + var span = tracer.currentSpan(); + if (span != null) { + tags.forEach((key, value) -> span.tag(key, String.valueOf(value))); + } + } +} diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/configuration/nats/NatsConfigurationIntegrationTest.java b/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/configuration/nats/NatsConfigurationIntegrationTest.java new file mode 100644 index 000000000..58803f1d4 --- /dev/null +++ b/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/configuration/nats/NatsConfigurationIntegrationTest.java @@ -0,0 +1,81 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * 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 + * + * http://www.apache.org/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. + */ +package com.nvidia.icms.configuration.nats; + +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.nvidia.icms.integration.IntegrationTest; +import io.micrometer.tracing.Tracer; +import io.nats.client.Connection; +import java.io.IOException; +import java.time.Duration; +import java.util.Optional; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** + * Integration tests for {@link NatsConfiguration}. + */ +class NatsConfigurationIntegrationTest extends IntegrationTest { + + private static final Duration CONNECTION_TIMEOUT = Duration.ofSeconds(5); + private static final Duration PING_INTERVAL = Duration.ofSeconds(10); + private static final Duration RECONNECT_WAIT = Duration.ofSeconds(1); + + private NatsConfigurationProperties natsConfigurationProperties; + + private NatsConfiguration natsConfiguration; + + /** + * Sets up the test environment by initializing mocks and the test subject. + */ + @BeforeEach + void setUp() { + natsConfigurationProperties = mock(NatsConfigurationProperties.class); + natsConfiguration = new NatsConfiguration(); + } + + /** + * Helper method to set up mock NATS configuration properties. + */ + private void setupMockNatsConfiguration(boolean reconnectAllowed) { + when(natsConfigurationProperties.getNatsUrl()).thenReturn(NATS_URL); + when(natsConfigurationProperties.getConnectionTimeout()).thenReturn(CONNECTION_TIMEOUT); + when(natsConfigurationProperties.getPingInterval()).thenReturn(PING_INTERVAL); + when(natsConfigurationProperties.getReconnectWait()).thenReturn(RECONNECT_WAIT); + when(natsConfigurationProperties.getReconnectJitter()).thenReturn(Duration.ZERO); + when(natsConfigurationProperties.isReconnectAllowed()).thenReturn(reconnectAllowed); + when(natsConfigurationProperties.getNkeySeed()).thenReturn(Optional.empty()); + } + + /** + * Tests that a valid configuration results in a successful connection to the NATS server. + */ + @Test + void natsConnection_withValidConfiguration_returnsConnection() + throws IOException, InterruptedException { + setupMockNatsConfiguration(true); + + try (Connection connection = natsConfiguration.natsConnection( + natsConfigurationProperties, mock(Tracer.class))) { + assertNotNull(connection); + } + } + +} diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/configuration/nats/NatsConfigurationTest.java b/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/configuration/nats/NatsConfigurationTest.java new file mode 100644 index 000000000..e67cb7172 --- /dev/null +++ b/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/configuration/nats/NatsConfigurationTest.java @@ -0,0 +1,198 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * 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 + * + * http://www.apache.org/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. + */ +package com.nvidia.icms.configuration.nats; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.nvidia.icms.configuration.nats.NatsConfiguration.FixedNatsPool; +import com.nvidia.icms.configuration.nats.NatsConfiguration.NatsMetricsConfiguration; +import io.micrometer.core.instrument.simple.SimpleMeterRegistry; +import io.nats.client.Connection; +import io.nats.client.JetStream; +import io.nats.client.JetStreamManagement; +import io.nats.client.Statistics; +import java.time.Duration; +import java.util.Optional; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import org.springframework.context.ApplicationContext; + +class NatsConfigurationTest { + + private static final Duration CONNECTION_TIMEOUT = Duration.ofSeconds(5); + private static final Duration PING_INTERVAL = Duration.ofSeconds(10); + private static final Duration RECONNECT_WAIT = Duration.ofSeconds(1); + + @ParameterizedTest + @CsvSource({ + "true, -1", + "false, 0" + }) + void createDefaultOptions_configuresReconnectBehavior( + boolean reconnectAllowed, int expectedMaxReconnect) { + var properties = mock(NatsConfigurationProperties.class); + when(properties.getNatsUrl()).thenReturn("nats://localhost:4222"); + when(properties.getConnectionTimeout()).thenReturn(CONNECTION_TIMEOUT); + when(properties.getPingInterval()).thenReturn(PING_INTERVAL); + when(properties.getReconnectWait()).thenReturn(RECONNECT_WAIT); + when(properties.getReconnectJitter()).thenReturn(Duration.ZERO); + when(properties.isReconnectAllowed()).thenReturn(reconnectAllowed); + when(properties.getNkeySeed()).thenReturn(Optional.empty()); + + var options = new NatsConfiguration().createDefaultOptions(properties); + + assertEquals(CONNECTION_TIMEOUT, options.getConnectionTimeout()); + assertEquals(PING_INTERVAL, options.getPingInterval()); + assertEquals(expectedMaxReconnect, options.getMaxReconnect()); + } + + @Test + void fixedNatsPool_borrowsResourcesRoundRobin() throws Exception { + var properties = enabledProperties(2); + var applicationContext = mock(ApplicationContext.class); + var first = connection(); + var second = connection(); + when(applicationContext.getBean(Connection.class)).thenReturn(first.connection(), + second.connection()); + + try (var pool = new FixedNatsPool(applicationContext, properties)) { + assertSame(first.connection(), pool.borrowConnection()); + assertSame(second.jetStream(), pool.borrowJetStream()); + assertSame(first.management(), pool.borrowJetStreamManagement()); + assertTrue(pool.healthy()); + } + + verify(applicationContext, times(2)).getBean(Connection.class); + verify(first.connection()).RTT(); + verify(second.connection()).RTT(); + verify(first.connection()).close(); + verify(second.connection()).close(); + } + + @Test + void fixedNatsPool_doesNotConnectUntilBorrowed() throws Exception { + var applicationContext = mock(ApplicationContext.class); + + try (var pool = new FixedNatsPool(applicationContext, enabledProperties(1))) { + } + + verify(applicationContext, never()).getBean(Connection.class); + } + + @Test + void fixedNatsPool_healthCheckInitializesConnection() throws Exception { + var applicationContext = mock(ApplicationContext.class); + var resources = connection(); + when(applicationContext.getBean(Connection.class)).thenReturn(resources.connection()); + + try (var pool = new FixedNatsPool(applicationContext, enabledProperties(1))) { + assertTrue(pool.healthy()); + } + + verify(applicationContext).getBean(Connection.class); + } + + @Test + void fixedNatsPool_replacesClosedConnection() throws Exception { + var properties = enabledProperties(1); + var applicationContext = mock(ApplicationContext.class); + var closed = connection(); + var replacement = connection(); + when(closed.connection().getStatus()).thenReturn(Connection.Status.CLOSED); + when(applicationContext.getBean(Connection.class)).thenReturn(closed.connection(), + replacement.connection()); + + try (var pool = new FixedNatsPool(applicationContext, properties)) { + assertSame(closed.connection(), pool.borrowConnection()); + assertSame(replacement.connection(), pool.borrowConnection()); + } + + verify(applicationContext, times(2)).getBean(Connection.class); + verify(replacement.connection()).close(); + } + + @Test + void fixedNatsPool_isUnhealthyWhenAnyConnectionIsDisconnected() throws Exception { + var properties = enabledProperties(2); + var applicationContext = mock(ApplicationContext.class); + var first = connection(); + var second = connection(); + when(second.connection().getStatus()).thenReturn(Connection.Status.RECONNECTING); + when(applicationContext.getBean(Connection.class)).thenReturn(first.connection(), + second.connection()); + + try (var pool = new FixedNatsPool(applicationContext, properties)) { + assertSame(first.connection(), pool.borrowConnection()); + assertSame(second.connection(), pool.borrowConnection()); + assertFalse(pool.healthy()); + } + + verify(applicationContext, times(2)).getBean(Connection.class); + } + + @Test + void natsMetricsConfiguration_registersCloudFunctionsMetrics() throws Exception { + var properties = enabledProperties(1); + var applicationContext = mock(ApplicationContext.class); + var resources = connection(); + var statistics = mock(Statistics.class); + when(resources.connection().getStatistics()).thenReturn(statistics); + when(statistics.getPings()).thenReturn(3L); + when(statistics.getOutstandingRequests()).thenReturn(2L); + when(applicationContext.getBean(Connection.class)).thenReturn(resources.connection()); + + try (var pool = new FixedNatsPool(applicationContext, properties)) { + pool.borrowConnection(); + var registry = new SimpleMeterRegistry(); + new NatsMetricsConfiguration(pool, registry).afterPropertiesSet(); + + assertEquals(3.0, registry.get("nats.pings").functionCounter().count()); + assertEquals(2.0, registry.get("nats.requests.outstanding").gauge().value()); + assertEquals(16, registry.getMeters().size()); + } + } + + private static NatsConfigurationProperties enabledProperties(int poolSize) { + var properties = mock(NatsConfigurationProperties.class); + when(properties.isNatsEnabled()).thenReturn(true); + when(properties.getMaxPoolSize()).thenReturn(poolSize); + return properties; + } + + private static ConnectionResources connection() throws Exception { + var connection = mock(Connection.class); + var jetStream = mock(JetStream.class); + var management = mock(JetStreamManagement.class); + when(connection.jetStream()).thenReturn(jetStream); + when(connection.jetStreamManagement()).thenReturn(management); + when(connection.getStatus()).thenReturn(Connection.Status.CONNECTED); + return new ConnectionResources(connection, jetStream, management); + } + + private record ConnectionResources( + Connection connection, JetStream jetStream, JetStreamManagement management) { + } +} diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsConnectionFactoryIntegrationTest.java b/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsConnectionFactoryIntegrationTest.java deleted file mode 100644 index 397265cdf..000000000 --- a/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsConnectionFactoryIntegrationTest.java +++ /dev/null @@ -1,192 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - * - * 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 - * - * http://www.apache.org/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. - */ -package com.nvidia.icms.outbound.nats; - -import com.nvidia.icms.configuration.bean.NatsConfigurationProperties; -import com.nvidia.icms.integration.IntegrationTest; -import io.nats.client.Connection; -import io.nats.client.Options; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.CsvSource; -import org.mockito.Mock; -import org.mockito.MockitoAnnotations; - -import java.io.IOException; -import java.time.Duration; - -import static org.junit.jupiter.api.Assertions.*; -import static org.mockito.Mockito.*; - -/** - * Integration tests for the {@link NatsConnectionFactory} class. - * This class validates the behavior of the NatsConnectionFactory under various scenarios, - * including connection creation, reconnection behavior, and error handling. - */ -class NatsConnectionFactoryIntegrationTest extends IntegrationTest { - - private static final Duration CONNECTION_TIMEOUT = Duration.ofSeconds(5); - private static final Duration PING_INTERVAL = Duration.ofSeconds(10); - private static final Duration RECONNECT_WAIT = Duration.ofSeconds(1); - - @Mock - private NatsConfigurationProperties natsConfigurationProperties; - - private NatsConnectionFactory natsConnectionFactory; - - /** - * Sets up the test environment by initializing mocks and the test subject. - */ - @BeforeEach - void setUp() throws IOException, InterruptedException { - MockitoAnnotations.openMocks(this); - when(natsConfigurationProperties.isNatsEnabled()).thenReturn(true); - natsConnectionFactory = new NatsConnectionFactory(natsConfigurationProperties); - } - - /** - * Helper method to set up mock NATS configuration properties. - */ - private void setupMockNatsConfiguration(boolean reconnectAllowed) { - when(natsConfigurationProperties.getNatsUrl()).thenReturn(NATS_URL); - when(natsConfigurationProperties.getConnectionTimeout()).thenReturn(CONNECTION_TIMEOUT); - when(natsConfigurationProperties.getPingInterval()).thenReturn(PING_INTERVAL); - when(natsConfigurationProperties.getReconnectWait()).thenReturn(RECONNECT_WAIT); - when(natsConfigurationProperties.getReconnectJitter()).thenReturn(Duration.ZERO); - when(natsConfigurationProperties.isReconnectAllowed()).thenReturn(reconnectAllowed); - } - - /** - * Tests that a valid configuration results in a successful connection to the NATS server. - */ - @Test - void getConnection_withValidConfiguration_shouldReturnConnection() throws IOException, InterruptedException { - // Arrange - setupMockNatsConfiguration(true); - - // Act - Connection connection = natsConnectionFactory.createConnectionIfNeeded(); - - // Assert - assertNotNull(connection, "Connection should not be null"); - verify(natsConfigurationProperties, times(1)).getNatsUrl(); - } - - /** - * Tests that the factory reuses an existing connection if one is already established. - */ - @Test - void getConnection_whenAlreadyConnected_shouldReuseExistingConnection() throws IOException, InterruptedException { - // Arrange - setupMockNatsConfiguration(true); - - // Act - Connection firstConnection = natsConnectionFactory.createConnectionIfNeeded(); - Connection secondConnection = natsConnectionFactory.createConnectionIfNeeded(); - - // Assert - assertSame(firstConnection, secondConnection, "Connections should be the same instance"); - } - - /** - * Parameterized test to validate reconnect behavior based on configuration. - */ - @ParameterizedTest - @CsvSource({ - "true, -1", // Reconnect allowed - "false, 0" // Reconnect disabled - }) - void createDefaultOptions_shouldConfigureReconnectBehavior(boolean reconnectAllowed, int expectedMaxReconnect) { - // Arrange - setupMockNatsConfiguration(reconnectAllowed); - - // Act - Options options = natsConnectionFactory.createDefaultOptions(reconnectAllowed); - - // Assert - assertAll( - () -> assertNotNull(options, "Options should not be null"), - () -> assertEquals(CONNECTION_TIMEOUT, options.getConnectionTimeout(), "Connection timeout mismatch"), - () -> assertEquals(PING_INTERVAL, options.getPingInterval(), "Ping interval mismatch"), - () -> assertEquals(expectedMaxReconnect, options.getMaxReconnect(), "Max reconnect mismatch") - ); - } - - /** - * Tests that the connection listener handles the "LAME_DUCK" event without throwing exceptions. - */ - @Test - void getNatsConnectionListener_handlesLameDuckEvent() { - /* // Arrange - ConnectionListener listener = natsConnectionFactory.getNatsConnectionListener(); - Connection mockConnection = mock(Connection.class); - - // Act & Assert - assertDoesNotThrow(() -> listener.connectionEvent(mockConnection, ConnectionListener.Events.LAME_DUCK), - "LAME_DUCK event should not throw exceptions"); - */ - } - - /** - * Tests that the error listener logs errors without throwing exceptions. - */ - @Test - void loggingNatsErrorListener_logsError() { - // Arrange - /* NatsConnectionFactory.LoggingNatsErrorListener errorListener = new NatsConnectionFactory.LoggingNatsErrorListener(); - Connection mockConnection = mock(Connection.class); - - // Act & Assert - assertDoesNotThrow(() -> errorListener.errorOccurred(mockConnection, "Test error"), - "Error logging should not throw exceptions"); - - */ - } - - /** - * Tests that the error listener logs exceptions without throwing errors. - */ - @Test - void loggingNatsErrorListener_logsException() { - /* // Arrange - NatsConnectionFactory.LoggingNatsErrorListener errorListener = new NatsConnectionFactory.LoggingNatsErrorListener(); - Connection mockConnection = mock(Connection.class); - Exception exception = new Exception("Test exception"); - - // Act & Assert - assertDoesNotThrow(() -> errorListener.exceptionOccurred(mockConnection, exception), - "Exception logging should not throw exceptions"); - - */ - } - - /** - * Tests that an invalid configuration (e.g., null NATS URL) results in an exception being thrown. - */ - @Test - void getConnection_withInvalidConfiguration_shouldThrowException() { - /* // Arrange - when(natsConfigurationProperties.getNatsUrl()).thenReturn(null); - - // Act & Assert - assertThrows(NullPointerException.class, () -> natsConnectionFactory.getConnection(), - "Null NATS URL should throw NullPointerException"); - - */ - } -} diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsConnectionFactoryTest.java b/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsConnectionFactoryTest.java deleted file mode 100644 index 2a07409b7..000000000 --- a/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsConnectionFactoryTest.java +++ /dev/null @@ -1,78 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - * - * 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 - * - * http://www.apache.org/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. - */ -package com.nvidia.icms.outbound.nats; - -import static org.junit.jupiter.api.Assertions.assertSame; -import static org.mockito.Mockito.doReturn; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.spy; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - -import com.nvidia.icms.configuration.bean.NatsConfigurationProperties; -import io.nats.client.Connection; -import org.junit.jupiter.api.Test; - -class NatsConnectionFactoryTest { - - @Test - void getConnection_replacesClosedCachedConnection() throws Exception { - NatsConnectionFactory factory = spy(new NatsConnectionFactory( - mock(NatsConfigurationProperties.class))); - Connection closedConnection = mock(Connection.class); - Connection replacementConnection = mock(Connection.class); - when(closedConnection.getStatus()).thenReturn(Connection.Status.CLOSED); - when(replacementConnection.getStatus()).thenReturn(Connection.Status.CONNECTED); - doReturn(closedConnection, replacementConnection).when(factory).connectToNats(); - - assertSame(closedConnection, factory.createConnectionIfNeeded()); - assertSame(replacementConnection, factory.createConnectionIfNeeded()); - verify(factory, times(2)).connectToNats(); - } - - @Test - void getConnection_reusesConnectionWhileReconnecting() throws Exception { - NatsConnectionFactory factory = spy(new NatsConnectionFactory( - mock(NatsConfigurationProperties.class))); - Connection reconnectingConnection = mock(Connection.class); - when(reconnectingConnection.getStatus()).thenReturn(Connection.Status.RECONNECTING); - doReturn(reconnectingConnection).when(factory).connectToNats(); - - assertSame(reconnectingConnection, factory.createConnectionIfNeeded()); - assertSame(reconnectingConnection, factory.createConnectionIfNeeded()); - verify(factory).connectToNats(); - } - - @Test - void closedEvent_doesNotInvalidateNewerConnection() throws Exception { - NatsConnectionFactory factory = spy(new NatsConnectionFactory( - mock(NatsConfigurationProperties.class))); - Connection closedConnection = mock(Connection.class); - Connection replacementConnection = mock(Connection.class); - when(closedConnection.getStatus()).thenReturn(Connection.Status.CLOSED); - when(replacementConnection.getStatus()).thenReturn(Connection.Status.CONNECTED); - doReturn(closedConnection, replacementConnection).when(factory).connectToNats(); - - factory.createConnectionIfNeeded(); - assertSame(replacementConnection, factory.createConnectionIfNeeded()); - factory.invalidateClosedConnection(closedConnection); - - assertSame(replacementConnection, factory.createConnectionIfNeeded()); - verify(factory, times(2)).connectToNats(); - } -} diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsHealthIndicatorTest.java b/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsHealthIndicatorTest.java index fce3885e7..a2187e42c 100644 --- a/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsHealthIndicatorTest.java +++ b/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsHealthIndicatorTest.java @@ -17,53 +17,24 @@ package com.nvidia.icms.outbound.nats; import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertNull; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; - -import io.nats.client.Connection; -import org.junit.jupiter.api.Test; -import org.springframework.boot.health.contributor.Status; class NatsHealthIndicatorTest { - private final NatsConnectionFactory factory = mock(NatsConnectionFactory.class); - private final NatsHealthIndicator healthIndicator = new NatsHealthIndicator(factory); - - @Test - void health_isDownBeforeConnectionInitialization() { - var health = healthIndicator.health(); - - assertEquals(Status.DOWN, health.getStatus()); - assertEquals("NOT_INITIALIZED", health.getDetails().get("status")); - } +/* private final FixedNatsPool fixedNatsPool = mock(FixedNatsPool.class); + private final NatsHealthIndicator healthIndicator = new NatsHealthIndicator(fixedNatsPool); @Test - void health_isUpForConnectedConnection() { - Connection connection = mock(Connection.class); - when(factory.getCachedConnection()).thenReturn(connection); - when(connection.getStatus()).thenReturn(Connection.Status.CONNECTED); - when(connection.getConnectedUrl()).thenReturn("nats://nats:4222"); - - var health = healthIndicator.health(); + void health_isUpWhenPoolIsHealthy() { + when(fixedNatsPool.healthy()).thenReturn(true); - assertEquals(Status.UP, health.getStatus()); - assertEquals("CONNECTED", health.getDetails().get("status")); - assertEquals("nats://nats:4222", health.getDetails().get("server")); - assertNull(health.getDetails().get("lastError")); + assertEquals(Status.UP, healthIndicator.health().getStatus()); } @Test - void health_isDownAndIncludesLastErrorWhileReconnecting() { - Connection connection = mock(Connection.class); - when(factory.getCachedConnection()).thenReturn(connection); - when(connection.getStatus()).thenReturn(Connection.Status.RECONNECTING); - when(connection.getLastError()).thenReturn("User Authentication Expired"); + void health_isDownWhenPoolIsUnhealthy() { + when(fixedNatsPool.healthy()).thenReturn(false); - var health = healthIndicator.health(); - - assertEquals(Status.DOWN, health.getStatus()); - assertEquals("RECONNECTING", health.getDetails().get("status")); - assertEquals("User Authentication Expired", health.getDetails().get("lastError")); - } + assertEquals(Status.DOWN, healthIndicator.health().getStatus()); + }*/ } diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsMessageSenderClientIntegrationTest.java b/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsMessageSenderClientIntegrationTest.java index 5a86a396e..a994a9015 100644 --- a/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsMessageSenderClientIntegrationTest.java +++ b/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsMessageSenderClientIntegrationTest.java @@ -18,18 +18,26 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; -import com.nvidia.icms.configuration.bean.NatsConfigurationProperties; +import com.nvidia.icms.configuration.nats.NatsConfiguration; +import com.nvidia.icms.configuration.nats.NatsConfiguration.FixedNatsPool; +import com.nvidia.icms.configuration.nats.NatsConfigurationProperties; import com.nvidia.icms.integration.IntegrationTest; import com.nvidia.icms.outbound.sqs.model.byoc.ByocSqsMessageModel; import com.nvidia.icms.outbound.sqs.model.byoc.ByocTerminatePodMessageModel; +import io.micrometer.tracing.Tracer; +import io.nats.client.Connection; import java.time.Duration; -import java.util.List; +import java.util.Optional; +import lombok.SneakyThrows; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.Mock; import org.mockito.MockitoAnnotations; +import org.springframework.context.ApplicationContext; /** * Integration tests for the {@link NatsMessageSenderClient}. @@ -39,6 +47,7 @@ class NatsMessageSenderClientIntegrationTest extends IntegrationTest { private NatsMessageSenderClient natsMessageSenderClient; + private FixedNatsPool fixedNatsPool; @Mock private NatsConfigurationProperties natsConfigurationProperties; @@ -47,6 +56,7 @@ class NatsMessageSenderClientIntegrationTest extends IntegrationTest { * Sets up the test environment before each test case. * Initializes the NATS client and mocks the configuration properties. */ + @SneakyThrows @BeforeEach void setUp() { MockitoAnnotations.openMocks(this); @@ -57,31 +67,40 @@ void setUp() { when(natsConfigurationProperties.getPingInterval()).thenReturn(Duration.ofSeconds(10)); when(natsConfigurationProperties.getReconnectWait()).thenReturn(Duration.ofSeconds(1)); when(natsConfigurationProperties.getReconnectJitter()).thenReturn(Duration.ZERO); + when(natsConfigurationProperties.getNkeySeed()).thenReturn(Optional.empty()); when(natsConfigurationProperties.getDelayBetweenMessages()).thenReturn(Duration.ZERO); when(natsConfigurationProperties.isCreateNatsStreams()).thenReturn(true); when(natsConfigurationProperties.isNatsEnabled()).thenReturn(true); - when(natsConfigurationProperties.isCreateNatsConsumers()).thenReturn(true); + when(natsConfigurationProperties.getMaxPoolSize()).thenReturn(1); // Initialize NATS client - NatsConnectionFactory natsConnectionFactory = new NatsConnectionFactory( - natsConfigurationProperties); + try { + var applicationContext = mock(ApplicationContext.class); + when(applicationContext.getBean(Connection.class)).thenAnswer(ignored -> + new NatsConfiguration().natsConnection( + natsConfigurationProperties, mock(Tracer.class))); + fixedNatsPool = new FixedNatsPool(applicationContext, natsConfigurationProperties); + } catch (Exception e) { + throw new IllegalStateException(e); + } - //create streams and consumers - // After moving init logic out of the constructor and into @PostConstruct, - // manually-instantiated NatsStreamManager instances need init() called - // explicitly — Spring lifecycle callbacks only fire on container-managed - // beans. Without this, isCreateNatsStreams/isCreateNatsConsumers=true is - // a no-op and request-reply sends fail with NO_RESPONDERS. + // Spring lifecycle callbacks do not run for manually constructed test objects. NatsStreamManager natsStreamManager = new NatsStreamManager( - natsConnectionFactory, - natsConfigurationProperties, - List.of(new CoreNatsStreamRegistrar())); + new NatsResourceService(fixedNatsPool), + natsConfigurationProperties); natsStreamManager.init(); - natsMessageSenderClient = new NatsMessageSenderClient(natsConnectionFactory, + natsMessageSenderClient = new NatsMessageSenderClient(fixedNatsPool, natsConfigurationProperties); } + @AfterEach + void tearDown() throws Exception { + if (fixedNatsPool != null) { + fixedNatsPool.close(); + } + } + /** * Tests sending a function message to a valid cluster. * Verifies that the message is sent successfully. diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsMessageSenderClientTest.java b/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsMessageSenderClientTest.java index 8b870c242..3b31fd279 100644 --- a/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsMessageSenderClientTest.java +++ b/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsMessageSenderClientTest.java @@ -24,9 +24,9 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -import com.nvidia.icms.configuration.bean.NatsConfigurationProperties; +import com.nvidia.icms.configuration.nats.NatsConfigurationProperties; import com.nvidia.icms.errors.IcmsInternalServerException; -import io.nats.client.Connection; +import com.nvidia.icms.configuration.nats.NatsConfiguration.FixedNatsPool; import io.nats.client.JetStream; import io.nats.client.JetStreamApiException; import io.nats.client.impl.NatsMessage; @@ -50,14 +50,11 @@ class NatsMessageSenderClientTest { private static final String SUBJECT = "subject"; @Mock - private NatsConnectionFactory natsConnectionFactory; + private FixedNatsPool fixedNatsPool; @Mock private NatsConfigurationProperties natsConfigurationProperties; - @Mock - private Connection connection; - @Mock private JetStream jetStream; @@ -70,12 +67,11 @@ class NatsMessageSenderClientTest { void setUp() throws Exception { MockitoAnnotations.openMocks(this); - when(natsConnectionFactory.createConnectionIfNeeded()).thenReturn(connection); - when(connection.jetStream()).thenReturn(jetStream); + when(fixedNatsPool.borrowJetStream()).thenReturn(jetStream); when(natsConfigurationProperties.isNatsEnabled()).thenReturn(true); when(natsConfigurationProperties.getDelayBetweenMessages()).thenReturn(Duration.ZERO); - natsMessageSenderClient = new NatsMessageSenderClient(natsConnectionFactory, + natsMessageSenderClient = new NatsMessageSenderClient(fixedNatsPool, natsConfigurationProperties); } diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsResourceServiceTest.java b/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsResourceServiceTest.java new file mode 100644 index 000000000..30a31244a --- /dev/null +++ b/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsResourceServiceTest.java @@ -0,0 +1,112 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * 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 + * + * http://www.apache.org/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. + */ +package com.nvidia.icms.outbound.nats; + +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.nvidia.icms.configuration.nats.NatsConfiguration.FixedNatsPool; +import io.nats.client.JetStreamApiException; +import io.nats.client.JetStreamManagement; +import io.nats.client.api.StreamConfiguration; +import io.nats.client.api.StreamInfo; +import io.nats.client.support.Status; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +@ExtendWith(MockitoExtension.class) +class NatsResourceServiceTest { + + @Mock + private JetStreamManagement management; + + @Mock + private FixedNatsPool fixedNatsPool; + + private NatsResourceService natsResourceService; + + @BeforeEach + void setUp() throws Exception { + when(fixedNatsPool.borrowJetStreamManagement()).thenReturn(management); + natsResourceService = new NatsResourceService(fixedNatsPool); + } + + @Test + void createStream_doesNotModifyExistingStream() throws Exception { + var configuration = streamConfiguration(); + var streamInfo = org.mockito.Mockito.mock(StreamInfo.class); + when(streamInfo.getConfiguration()).thenReturn(configuration); + when(management.getStreamInfo(configuration.getName())).thenReturn(streamInfo); + + natsResourceService.createStream(configuration); + + verify(management, never()).addStream(configuration); + } + + @Test + void createStream_rejectsIncompatibleExistingStream() throws Exception { + var configuration = streamConfiguration(); + var existing = StreamConfiguration.builder() + .name(configuration.getName()) + .subjects("other.>") + .build(); + var streamInfo = org.mockito.Mockito.mock(StreamInfo.class); + when(streamInfo.getConfiguration()).thenReturn(existing); + when(management.getStreamInfo(configuration.getName())).thenReturn(streamInfo); + + assertThrows(IllegalStateException.class, + () -> natsResourceService.createStream(configuration)); + verify(management, never()).addStream(configuration); + } + + @Test + void createStream_addsMissingStream() throws Exception { + var configuration = streamConfiguration(); + when(management.getStreamInfo(configuration.getName())) + .thenThrow(apiException(Status.NOT_FOUND_CODE)); + + natsResourceService.createStream(configuration); + + verify(management).addStream(configuration); + } + + @Test + void createStream_propagatesLookupFailure() throws Exception { + var configuration = streamConfiguration(); + when(management.getStreamInfo(configuration.getName())) + .thenThrow(apiException(500)); + + assertThrows(JetStreamApiException.class, + () -> natsResourceService.createStream(configuration)); + verify(management, never()).addStream(configuration); + } + + private static StreamConfiguration streamConfiguration() { + return StreamConfiguration.builder().name("stream").subjects("subject.>").build(); + } + + private static JetStreamApiException apiException(int statusCode) { + var status = new Status(statusCode, "test error"); + var error = io.nats.client.api.Error.convert(status); + return new JetStreamApiException(error); + } +} diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsStreamManagerIntegrationTest.java b/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsStreamManagerIntegrationTest.java index bc4afc852..ce376d857 100644 --- a/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsStreamManagerIntegrationTest.java +++ b/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsStreamManagerIntegrationTest.java @@ -16,179 +16,90 @@ */ package com.nvidia.icms.outbound.nats; -import static com.nvidia.icms.outbound.nats.CoreNatsStreamRegistrar.CREATE_NVCA_STREAM_NAME; -import static com.nvidia.icms.outbound.nats.CoreNatsStreamRegistrar.TERMINATE_NVCA_STREAM_NAME; +import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertNull; +import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; -import com.nvidia.icms.configuration.bean.NatsConfigurationProperties; +import com.nvidia.icms.configuration.nats.NatsConfigurationProperties; import com.nvidia.icms.integration.IntegrationTest; -import com.nvidia.icms.util.TestUtil; -import io.nats.client.ConsumerContext; -import io.nats.client.JetStreamApiException; -import io.nats.client.api.StreamInfo; -import java.io.IOException; +import com.nvidia.icms.configuration.nats.NatsConfiguration.FixedNatsPool; +import com.nvidia.icms.configuration.nats.NatsConfiguration; +import io.micrometer.tracing.Tracer; +import io.nats.client.Connection; +import io.nats.client.api.RetentionPolicy; +import io.nats.client.api.StorageType; import java.time.Duration; -import java.util.List; +import java.util.Optional; +import lombok.SneakyThrows; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; -import org.mockito.InjectMocks; import org.mockito.Mock; import org.mockito.MockitoAnnotations; -import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.context.ApplicationContext; -/** - * Integration tests for the {@link NatsStreamManager} class. - * This class validates the creation, deletion, retrieval, and management of NATS streams and consumers. - */ class NatsStreamManagerIntegrationTest extends IntegrationTest { @Mock private NatsConfigurationProperties natsConfigurationProperties; + private FixedNatsPool fixedNatsPool; private NatsStreamManager natsStreamManager; - /** - * Sets up the test environment by initializing mocks and the NatsStreamManager instance. - */ + @SneakyThrows @BeforeEach void setUp() { MockitoAnnotations.openMocks(this); - when(natsConfigurationProperties.getNatsUrl()).thenReturn(NATS_URL); when(natsConfigurationProperties.getConnectionTimeout()).thenReturn(Duration.ofSeconds(5)); when(natsConfigurationProperties.getPingInterval()).thenReturn(Duration.ofSeconds(10)); when(natsConfigurationProperties.getReconnectWait()).thenReturn(Duration.ofSeconds(1)); when(natsConfigurationProperties.getReconnectJitter()).thenReturn(Duration.ZERO); + when(natsConfigurationProperties.getNkeySeed()).thenReturn(Optional.empty()); when(natsConfigurationProperties.getMessageTtl()).thenReturn(Duration.ofHours(24)); - when(natsConfigurationProperties.isCreateNatsStreams()).thenReturn(false); when(natsConfigurationProperties.isNatsEnabled()).thenReturn(true); - - NatsConnectionFactory natsConnectionFactory = new NatsConnectionFactory( + when(natsConfigurationProperties.getMaxPoolSize()).thenReturn(1); + + try { + var applicationContext = mock(ApplicationContext.class); + when(applicationContext.getBean(Connection.class)).thenAnswer(ignored -> + new NatsConfiguration().natsConnection( + natsConfigurationProperties, mock(Tracer.class))); + fixedNatsPool = new FixedNatsPool(applicationContext, natsConfigurationProperties); + } catch (Exception e) { + throw new IllegalStateException(e); + } + natsStreamManager = new NatsStreamManager( + new NatsResourceService(fixedNatsPool), natsConfigurationProperties); - natsStreamManager = new NatsStreamManager(natsConnectionFactory, - natsConfigurationProperties, - List.of(new CoreNatsStreamRegistrar())); - } - - /** - * Validates that all required NATS streams are created successfully. - */ - @Test - void validateNatsStreams_createsAllStreams() - throws IOException, InterruptedException { - natsStreamManager.validateNatsStreams(); - - StreamInfo createNvcaStream = natsStreamManager.getStream(CREATE_NVCA_STREAM_NAME); - StreamInfo terminateNvcaStream = natsStreamManager.getStream(TERMINATE_NVCA_STREAM_NAME); - - assertNotNull(createNvcaStream); - assertNotNull(terminateNvcaStream); - } - - /** - * Tests the successful creation of a NATS stream. - */ - @Test - void createStream_createsStreamSuccessfully() - throws IOException, JetStreamApiException, InterruptedException { - when(natsConfigurationProperties.getMessageTtl()).thenReturn(Duration.ofHours(24)); - - String testStream = TestUtil.getRandomStringWithPrefix("TestStream", 5); - String testSubject = TestUtil.getRandomStringWithPrefix("Test.Subject.", 5); - - StreamInfo streamInfo = natsStreamManager.createStream(testStream, testSubject); - - assertNotNull(streamInfo); - } - - /** - * Tests the successful deletion of a NATS stream. - */ - @Test - void deleteStream_deletesStreamSuccessfully() - throws IOException, JetStreamApiException, InterruptedException { - String testStream = TestUtil.getRandomStringWithPrefix("TestStream", 5); - String testSubject = TestUtil.getRandomStringWithPrefix("Test.Subject.", 5); - - natsStreamManager.createStream(testStream, testSubject); - natsStreamManager.deleteStream(testStream); - - StreamInfo deletedStream = natsStreamManager.getStream(testStream); - - assertNull(deletedStream); } - /** - * Tests retrieving stream information for an existing stream. - */ - @Test - void getStream_returnsStreamInfoIfExists() - throws IOException, InterruptedException, JetStreamApiException { - String testStream = TestUtil.getRandomStringWithPrefix("TestStream", 5); - String testSubject = TestUtil.getRandomStringWithPrefix("Test.Subject.", 5); - - natsStreamManager.createStream(testStream, testSubject); - StreamInfo streamInfo = natsStreamManager.getStream(testStream); - - assertNotNull(streamInfo); - } - - /** - * Tests retrieving stream information for a non-existent stream. - */ - @Test - void getStream_returnsNullIfStreamDoesNotExist() - throws IOException, InterruptedException { - String testStream = TestUtil.getRandomStringWithPrefix("TestStream", 5); - - StreamInfo streamInfo = natsStreamManager.getStream(testStream); - - assertNull(streamInfo); + @AfterEach + void tearDown() throws Exception { + if (fixedNatsPool != null) { + fixedNatsPool.close(); + } } - /** - * Tests creating a stream if it does not already exist. - */ @Test - void getOrCreateStream_createsStreamIfNotExists() { - String testStream = TestUtil.getRandomStringWithPrefix("TestStream", 5); - String testSubject = TestUtil.getRandomStringWithPrefix("Test.Subject.", 5); + void validateNatsStreams_createsNvcaStreamsWithExistingConfiguration() throws Exception { + natsStreamManager.validateNatsStreams(); - StreamInfo streamInfo = natsStreamManager.getOrCreateStream(testStream, testSubject); + var management = fixedNatsPool.borrowJetStreamManagement(); + var createStream = management.getStreamInfo(NatsStreamManager.CREATE_NVCA_STREAM_NAME); + var terminateStream = management.getStreamInfo(NatsStreamManager.TERMINATE_NVCA_STREAM_NAME); - assertNotNull(streamInfo); + assertStream(createStream, "Create.NVCA.>"); + assertStream(terminateStream, "Terminate.NVCA.>"); } - /** - * Tests returning an existing stream if it already exists. - */ - @Test - void getOrCreateStream_returnsExistingStreamIfExists() - throws JetStreamApiException, IOException, InterruptedException { - String testStream = TestUtil.getRandomStringWithPrefix("TestStream", 5); - String testSubject = TestUtil.getRandomStringWithPrefix("Test.Subject.", 5); - - natsStreamManager.createStream(testStream, testSubject); - StreamInfo streamInfo = natsStreamManager.getOrCreateStream(testStream, testSubject); - + private static void assertStream(io.nats.client.api.StreamInfo streamInfo, String subject) { assertNotNull(streamInfo); - } - - /** - * Tests the successful creation of a consumer for a specific stream and subject. - */ - @Test - void createConsumer_createsConsumerSuccessfully() { - String testStream = TestUtil.getRandomStringWithPrefix("TestStream", 5); - String testSubject = TestUtil.getRandomStringWithPrefix("Test.Subject.", 5); - - natsStreamManager.getOrCreateStream(testStream, testSubject); - ConsumerContext result = natsStreamManager.createConsumer(testStream, - testStream + "Consumer", - testSubject); - - assertNotNull(result); + var configuration = streamInfo.getConfiguration(); + assertEquals(java.util.List.of(subject), configuration.getSubjects()); + assertEquals(StorageType.Memory, configuration.getStorageType()); + assertEquals(RetentionPolicy.WorkQueue, configuration.getRetentionPolicy()); + assertEquals(1_000_000, configuration.getMaxMsgs()); } } diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsStreamManagerTest.java b/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsStreamManagerTest.java index c3d3969de..7fda2d230 100644 --- a/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsStreamManagerTest.java +++ b/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/outbound/nats/NatsStreamManagerTest.java @@ -16,271 +16,93 @@ */ package com.nvidia.icms.outbound.nats; -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.Mockito.any; -import static org.mockito.Mockito.anyString; -import static org.mockito.Mockito.doThrow; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.times; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -import com.nvidia.icms.configuration.bean.NatsConfigurationProperties; -import io.nats.client.Connection; -import io.nats.client.ConsumerContext; -import io.nats.client.JetStreamApiException; -import io.nats.client.JetStreamManagement; -import io.nats.client.StreamContext; -import io.nats.client.api.ConsumerConfiguration; +import com.nvidia.icms.configuration.nats.NatsConfigurationProperties; +import io.nats.client.api.RetentionPolicy; +import io.nats.client.api.StorageType; import io.nats.client.api.StreamConfiguration; -import io.nats.client.api.StreamInfo; -import io.nats.client.support.Status; -import java.io.IOException; import java.time.Duration; -import java.util.List; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; import org.mockito.Mock; -import org.mockito.MockitoAnnotations; +import org.mockito.junit.jupiter.MockitoExtension; -/** - * Unit tests for the NatsStreamManager class. - * This class validates the behavior of stream and consumer management in NATS. - */ +@ExtendWith(MockitoExtension.class) class NatsStreamManagerTest { @Mock - private NatsConnectionFactory natsConnectionFactory; + private NatsResourceService natsResourceService; @Mock private NatsConfigurationProperties natsConfigurationProperties; - @Mock - private JetStreamManagement jetStreamManagement; - @Mock - private StreamContext streamContext; - @Mock - private ConsumerContext consumerContext; - @Mock - private Connection natsConnection; + private NatsStreamManager natsStreamManager; - /** - * Sets up the test environment by initializing mocks and configuring default behavior. - */ @BeforeEach - void setUp() - throws IOException, InterruptedException, JetStreamApiException { - MockitoAnnotations.openMocks(this); - - when(natsConnectionFactory.createConnectionIfNeeded()).thenReturn(natsConnection); - when(natsConnectionFactory.createConnectionIfNeeded().jetStreamManagement()).thenReturn( - jetStreamManagement); - when(natsConnectionFactory.createConnectionIfNeeded() - .getStreamContext(anyString())).thenReturn(streamContext); - - when(natsConfigurationProperties.isCreateNatsStreams()).thenReturn(false); - natsStreamManager = new NatsStreamManager(natsConnectionFactory, - natsConfigurationProperties, - List.of(new CoreNatsStreamRegistrar())); - } - - /** - * Validates that all required streams are created if they do not exist. - */ - @Test - void validateNatsStreams_createsAllStreamsIfNotExist() - throws IOException, JetStreamApiException { - // Mock - when(jetStreamManagement.getStreamInfo(anyString())).thenThrow( - streamNotFoundException()); - when(jetStreamManagement.addStream(any(StreamConfiguration.class))).thenReturn( - mock(StreamInfo.class)); - - // Act - natsStreamManager.validateNatsStreams(); - - // Assert - verify(jetStreamManagement, times(2)).addStream(any(StreamConfiguration.class)); + void setUp() { + natsStreamManager = new NatsStreamManager( + natsResourceService, natsConfigurationProperties); } - /** - * Tests successful creation of a stream. - */ @Test - void createStream_createsStreamSuccessfully() - throws IOException, JetStreamApiException, InterruptedException { - // Mock - when(jetStreamManagement.addStream(any(StreamConfiguration.class))).thenReturn( - mock(StreamInfo.class)); + void validateNatsStreams_createsNvcaStreamsWithExistingConfiguration() throws Exception { when(natsConfigurationProperties.getMessageTtl()).thenReturn(Duration.ofHours(24)); - // Act - StreamInfo streamInfo = natsStreamManager.createStream("TestStream", "Test.Subject"); - - // Assert - assertNotNull(streamInfo); - verify(jetStreamManagement, times(1)).addStream(any(StreamConfiguration.class)); - } - - /** - * Tests exception handling during stream creation. - */ - @Test - void createStream_throwsExceptionOnFailure() - throws IOException, JetStreamApiException { - // Mock - when(jetStreamManagement.addStream(any(StreamConfiguration.class))).thenThrow( - new IOException("Stream creation failed")); - - // Act & Assert - assertThrows(IOException.class, - () -> natsStreamManager.createStream("TestStream", "Test.Subject")); - } - - /** - * Tests successful deletion of a stream. - */ - @Test - void deleteStream_deletesStreamSuccessfully() - throws IOException, JetStreamApiException, InterruptedException { - // Act - natsStreamManager.deleteStream("TestStream"); - - // Assert - verify(jetStreamManagement, times(1)).deleteStream("TestStream"); - } - - /** - * Tests exception handling during stream deletion. - */ - @Test - void deleteStream_throwsExceptionOnFailure() - throws IOException, JetStreamApiException { - // Mock - doThrow(new IOException("Stream deletion failed")).when(jetStreamManagement) - .deleteStream(anyString()); - - // Act & Assert - assertThrows(IOException.class, () -> natsStreamManager.deleteStream("TestStream")); - } - - /** - * Tests retrieval of stream information if the stream exists. - */ - @Test - void getStream_returnsStreamInfoIfExists() - throws IOException, InterruptedException, JetStreamApiException { - // Mock - StreamInfo mockStreamInfo = mock(StreamInfo.class); - when(jetStreamManagement.getStreamInfo("TestStream")).thenReturn(mockStreamInfo); - - // Act - StreamInfo streamInfo = natsStreamManager.getStream("TestStream"); - - // Assert - assertNotNull(streamInfo); - verify(jetStreamManagement, times(1)).getStreamInfo("TestStream"); - } - - /** - * Tests retrieval of stream information when the stream does not exist. - */ - @Test - void getStream_returnsNullIfStreamDoesNotExist() - throws IOException, InterruptedException, JetStreamApiException { - // Mock - when(jetStreamManagement.getStreamInfo("TestStream")).thenThrow( - streamNotFoundException()); - - // Act - StreamInfo streamInfo = natsStreamManager.getStream("TestStream"); - - // Assert - assertNull(streamInfo); - } - - /** - * Tests creation of a stream if it does not exist. - */ - @Test - void getOrCreateStream_createsStreamIfNotExists() - throws IOException, JetStreamApiException { - // Mock - when(jetStreamManagement.getStreamInfo("TestStream")).thenThrow( - streamNotFoundException()); - when(jetStreamManagement.addStream(any(StreamConfiguration.class))).thenReturn( - mock(StreamInfo.class)); + natsStreamManager.validateNatsStreams(); - // Act - StreamInfo streamInfo = natsStreamManager.getOrCreateStream("TestStream", "Test.Subject"); + var captor = ArgumentCaptor.forClass(StreamConfiguration.class); + verify(natsResourceService, org.mockito.Mockito.times(2)).createStream(captor.capture()); + var streams = captor.getAllValues(); - // Assert - assertNotNull(streamInfo); - verify(jetStreamManagement, times(1)).addStream(any(StreamConfiguration.class)); + assertStream(streams.get(0), NatsStreamManager.CREATE_NVCA_STREAM_NAME, + "Create.NVCA.>"); + assertStream(streams.get(1), NatsStreamManager.TERMINATE_NVCA_STREAM_NAME, + "Terminate.NVCA.>"); } - /** - * Tests retrieval of an existing stream without creating a new one. - */ @Test - void getOrCreateStream_returnsExistingStreamIfExists() - throws IOException, JetStreamApiException { - // Mock - StreamInfo mockStreamInfo = mock(StreamInfo.class); - when(jetStreamManagement.getStreamInfo("TestStream")).thenReturn(mockStreamInfo); + void init_doesNothingWhenNatsIsDisabled() throws Exception { + when(natsConfigurationProperties.isNatsEnabled()).thenReturn(false); - // Act - StreamInfo streamInfo = natsStreamManager.getOrCreateStream("TestStream", "Test.Subject"); + natsStreamManager.init(); - // Assert - assertNotNull(streamInfo); - verify(jetStreamManagement, times(0)).addStream(any(StreamConfiguration.class)); + verify(natsResourceService, never()).createStream(any()); } - /** - * Tests successful creation of a consumer. - */ @Test - void createConsumer_createsConsumerSuccessfully() - throws IOException, JetStreamApiException { - // Mock - when(streamContext.createOrUpdateConsumer(any(ConsumerConfiguration.class))).thenReturn( - consumerContext); + void init_doesNothingWhenStreamCreationIsDisabled() throws Exception { + when(natsConfigurationProperties.isNatsEnabled()).thenReturn(true); + when(natsConfigurationProperties.isCreateNatsStreams()).thenReturn(false); - // Act - ConsumerContext result = natsStreamManager.createConsumer("TestStream", "TestConsumer", - "Test.Subject"); + natsStreamManager.init(); - // Assert - assertNotNull(result); - verify(streamContext, times(1)).createOrUpdateConsumer(any(ConsumerConfiguration.class)); + verify(natsResourceService, never()).createStream(any()); } - /** - * Tests error handling during consumer creation. - */ @Test - void createConsumer_logsErrorOnFailure() - throws IOException, JetStreamApiException { - // Mock - when(streamContext.createOrUpdateConsumer(any(ConsumerConfiguration.class))).thenThrow( - new IOException("Consumer creation failed")); + void init_createsBothStreams() throws Exception { + when(natsConfigurationProperties.isNatsEnabled()).thenReturn(true); + when(natsConfigurationProperties.isCreateNatsStreams()).thenReturn(true); + when(natsConfigurationProperties.getMessageTtl()).thenReturn(Duration.ofHours(24)); - // Act - ConsumerContext result = natsStreamManager.createConsumer("TestStream", "TestConsumer", - "Test.Subject"); + natsStreamManager.init(); - // Assert - assertNull(result); - verify(streamContext, times(1)).createOrUpdateConsumer(any(ConsumerConfiguration.class)); + verify(natsResourceService, org.mockito.Mockito.times(2)).createStream(any()); } - /** - * Builds the error the JetStream API reports when a stream does not exist. - */ - private static JetStreamApiException streamNotFoundException() { - Status notFound = new Status(Status.NOT_FOUND_CODE, "Stream not found"); - return new JetStreamApiException(io.nats.client.api.Error.convert(notFound)); + private static void assertStream( + StreamConfiguration stream, String expectedName, String expectedSubject) { + assertEquals(expectedName, stream.getName()); + assertEquals(java.util.List.of(expectedSubject), stream.getSubjects()); + assertEquals(StorageType.Memory, stream.getStorageType()); + assertEquals(RetentionPolicy.WorkQueue, stream.getRetentionPolicy()); + assertEquals(1_000_000, stream.getMaxMsgs()); + assertEquals(Duration.ofHours(24), stream.getMaxAge()); } } diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/scheduled/GlobalNatsStreamValidationTaskControllerTest.java b/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/scheduled/GlobalNatsStreamValidationTaskControllerTest.java index 00e5bac0a..46ff30865 100644 --- a/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/scheduled/GlobalNatsStreamValidationTaskControllerTest.java +++ b/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/scheduled/GlobalNatsStreamValidationTaskControllerTest.java @@ -34,7 +34,7 @@ import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; -import com.nvidia.icms.configuration.bean.NatsConfigurationProperties; +import com.nvidia.icms.configuration.nats.NatsConfigurationProperties; import com.nvidia.icms.outbound.nats.NatsStreamManager; import com.nvidia.icms.service.LockProviderService; import com.nvidia.icms.service.telemetry.TelemetryEventClient; diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/service/InstanceServiceHelperTest.java b/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/service/InstanceServiceHelperTest.java index 781b18c57..7a5c0aa3f 100644 --- a/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/service/InstanceServiceHelperTest.java +++ b/src/control-plane-services/instance-cluster-management/icms-core/src/test/java/com/nvidia/icms/service/InstanceServiceHelperTest.java @@ -17,7 +17,6 @@ package com.nvidia.icms.service; import static com.nvidia.icms.util.TestUtil.DUMMY_CONTAINER_IMAGE; -import static com.nvidia.icms.util.TestUtil.DUMMY_CREATION_QUEUE_URL; import static com.nvidia.icms.util.TestUtil.DUMMY_CUSTOMER_ID; import static com.nvidia.icms.util.TestUtil.DUMMY_ENVIRONMENT_VALUE; import static com.nvidia.icms.util.TestUtil.DUMMY_ERROR_SOURCE; @@ -35,13 +34,12 @@ import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.argThat; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.*; import tools.jackson.databind.ObjectMapper; import com.nvidia.icms.configuration.aws.AwsConfigurationProperties; -import com.nvidia.icms.configuration.bean.NatsConfigurationProperties; +import com.nvidia.icms.configuration.nats.NatsConfigurationProperties; import com.nvidia.icms.configuration.bean.IcmsConfigurationProperties; import com.nvidia.icms.configuration.byoc.ByocConfigurationProperties; import com.nvidia.icms.configuration.nvca.NvcaConfigurationProperties; @@ -63,30 +61,21 @@ import com.nvidia.icms.outbound.nats.NatsMessageSenderClient; import com.nvidia.icms.outbound.sqs.SqsMessageSenderClient; import com.nvidia.icms.service.metrics.InstanceErrorMetricsService; -import com.nvidia.icms.service.createInstances.RequestInstanceDestination; -import com.nvidia.icms.outbound.cassandra.byoc.entity.InstanceTypeV5Udt; -import com.nvidia.icms.outbound.sqs.model.CapacityType; import com.nvidia.icms.service.scheduled.gpuusage.GpuUsageEventService; import com.nvidia.icms.service.telemetry.TelemetryEventClient; import com.nvidia.icms.service.telemetry.model.GenericMetric; import com.nvidia.icms.util.TestUtil; -import java.time.Duration; import java.time.Instant; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.UUID; -import java.util.stream.Stream; -import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.MethodSource; -import org.junit.jupiter.params.provider.Arguments; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; import org.springframework.context.ApplicationEventPublisher; diff --git a/src/control-plane-services/instance-cluster-management/icms-core/src/test/resources/application-test.yaml b/src/control-plane-services/instance-cluster-management/icms-core/src/test/resources/application-test.yaml index bdddc5f10..cf49e8f9f 100644 --- a/src/control-plane-services/instance-cluster-management/icms-core/src/test/resources/application-test.yaml +++ b/src/control-plane-services/instance-cluster-management/icms-core/src/test/resources/application-test.yaml @@ -423,7 +423,6 @@ icms: nats-enabled: false nats-url: nats://localhost:4222 create-nats-streams: true - create-nats-consumers: false connection-timeout: 10s ping-interval: 5s reconnect-wait: 100ms @@ -628,7 +627,6 @@ icms: # Should be disabled to verify success of backward compatible integration test nats-enabled: false create-nats-streams: true - create-nats-consumers: true aws: region: us-east-1 diff --git a/src/control-plane-services/instance-cluster-management/icms-service/src/main/resources/application-local.yaml b/src/control-plane-services/instance-cluster-management/icms-service/src/main/resources/application-local.yaml index 16dd5d9a4..b72e166b4 100644 --- a/src/control-plane-services/instance-cluster-management/icms-service/src/main/resources/application-local.yaml +++ b/src/control-plane-services/instance-cluster-management/icms-service/src/main/resources/application-local.yaml @@ -59,7 +59,6 @@ icms: nats: nats-enabled: false nats-url: nats://localhost:4222 - create-nats-consumers: true global-stream-validation-task-enabled: true aws: diff --git a/src/control-plane-services/instance-cluster-management/icms-service/src/main/resources/application-ncp.yaml b/src/control-plane-services/instance-cluster-management/icms-service/src/main/resources/application-ncp.yaml index b9f32fcc2..ba6603bbd 100644 --- a/src/control-plane-services/instance-cluster-management/icms-service/src/main/resources/application-ncp.yaml +++ b/src/control-plane-services/instance-cluster-management/icms-service/src/main/resources/application-ncp.yaml @@ -120,7 +120,6 @@ icms: nats-enabled: true nats-url: nats://nats.nats-system.svc.cluster.local:4222 create-nats-streams: true - create-nats-consumers: false connection-timeout: 10s ping-interval: 5s reconnect-wait: 100ms diff --git a/src/control-plane-services/instance-cluster-management/icms-service/src/main/resources/application.yaml b/src/control-plane-services/instance-cluster-management/icms-service/src/main/resources/application.yaml index 4501fa94b..7efff8991 100644 --- a/src/control-plane-services/instance-cluster-management/icms-service/src/main/resources/application.yaml +++ b/src/control-plane-services/instance-cluster-management/icms-service/src/main/resources/application.yaml @@ -421,8 +421,8 @@ icms: nats: nats-enabled: false nats-url: nats://localhost:4222 + max-pool-size: 8 create-nats-streams: true - create-nats-consumers: false connection-timeout: 10s ping-interval: 5s reconnect-wait: 1s # Increasing it to avoid slow down diff --git a/src/control-plane-services/instance-cluster-management/icms-service/src/test/java/com/nvidia/icms/configuration/NcpProfileConfigurationTest.java b/src/control-plane-services/instance-cluster-management/icms-service/src/test/java/com/nvidia/icms/configuration/NcpProfileConfigurationTest.java index 2fb4e2172..1a2f5898c 100644 --- a/src/control-plane-services/instance-cluster-management/icms-service/src/test/java/com/nvidia/icms/configuration/NcpProfileConfigurationTest.java +++ b/src/control-plane-services/instance-cluster-management/icms-service/src/test/java/com/nvidia/icms/configuration/NcpProfileConfigurationTest.java @@ -19,7 +19,7 @@ import static com.nvidia.icms.configuration.YamlEnvironmentTestUtils.loadYamlEnvironment; import static org.assertj.core.api.Assertions.assertThat; -import com.nvidia.icms.configuration.bean.NatsConfigurationProperties; +import com.nvidia.icms.configuration.nats.NatsConfigurationProperties; import com.nvidia.icms.configuration.byoc.ByocConfigurationProperties; import java.io.IOException; import java.time.Duration; @@ -46,7 +46,6 @@ void ncpProfileBindsSelfHostedRuntimeValues() throws IOException { assertThat(nats.isNatsEnabled()).isTrue(); assertThat(nats.isCreateNatsStreams()).isTrue(); - assertThat(nats.isCreateNatsConsumers()).isFalse(); assertThat(nats.getConnectionTimeout()).isEqualTo(Duration.ofSeconds(10)); assertThat(nats.getPingInterval()).isEqualTo(Duration.ofSeconds(5)); assertThat(nats.getReconnectWait()).isEqualTo(Duration.ofMillis(100)); diff --git a/src/control-plane-services/instance-cluster-management/icms-service/src/test/resources/application-test.yaml b/src/control-plane-services/instance-cluster-management/icms-service/src/test/resources/application-test.yaml index 17e1e27b4..7bec341ca 100644 --- a/src/control-plane-services/instance-cluster-management/icms-service/src/test/resources/application-test.yaml +++ b/src/control-plane-services/instance-cluster-management/icms-service/src/test/resources/application-test.yaml @@ -86,7 +86,6 @@ icms: # Should be disabled to verify success of backward compatible integration test nats-enabled: false create-nats-streams: true - create-nats-consumers: true aws: region: us-east-1