001/*
002 * Licensed to the Apache Software Foundation (ASF) under one or more
003 * contributor license agreements.  See the NOTICE file distributed with
004 * this work for additional information regarding copyright ownership.
005 * The ASF licenses this file to You under the Apache License, Version 2.0
006 * (the "License"); you may not use this file except in compliance with
007 * the License.  You may obtain a copy of the License at
008 *
009 *      http://www.apache.org/licenses/LICENSE-2.0
010 *
011 * Unless required by applicable law or agreed to in writing, software
012 * distributed under the License is distributed on an "AS IS" BASIS,
013 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
014 * See the License for the specific language governing permissions and
015 * limitations under the License.
016 */
017package org.apache.juneau.rest.server.health;
018
019import java.util.concurrent.*;
020
021import org.apache.juneau.commons.inject.*;
022
023/**
024 * Shared, observable readiness flag consumed by {@link HealthAggregator} to gate the readiness probe.
025 *
026 * <p>
027 * This is the building block for zero-downtime shutdown.  When a microservice begins shutting down it flips
028 * the readiness state {@link #markOutOfService() out of service} <i>before</i> the listening connector stops.
029 * The {@link HealthAggregator} consults this state for the {@link HealthProbe#READY READY} probe so
030 * {@code /readyz} returns {@code 503} (out of service) as soon as shutdown begins &mdash; this lets a load
031 * balancer / Kubernetes stop routing new traffic to the pod while in-flight requests drain.  The
032 * {@link HealthProbe#LIVE LIVE} probe ({@code /livez}) is intentionally <b>not</b> gated, so Kubernetes does
033 * not kill the pod mid-drain.
034 *
035 * <h5 class='section'>Resolution:</h5>
036 * <p>
037 * {@link HealthAggregator} and the embedded-server shutdown hooks resolve the state via
038 * {@link #resolve(BeanStore)}: a {@code ReadinessState} bean registered in the relevant {@link BeanStore} wins;
039 * otherwise the process-wide {@link #shared() shared} instance is used.  In the standalone Jetty/Tomcat
040 * microservice runtime, the embedded-server lifecycle component (e.g. {@code JettyServerComponent}) constructs
041 * one per-service instance on {@code start()} &mdash; reusing an app/test-supplied {@code @Bean ReadinessState}
042 * if one is already registered &mdash; and explicitly publishes that <b>same</b> instance into both the
043 * microservice's bean store and the auto-mounted health-probe servlet's own bean store.  The two sides then
044 * observe each other's flips without ever touching {@link #shared()}, so stopping one default-configured
045 * microservice does not affect any other microservice's readiness in the same JVM.  {@link #shared()} remains
046 * the last-resort fallback for embeddings that skip that dual-store publish (a plain servlet container, Spring
047 * Boot without this integration, or unit tests that never register a bean).
048 *
049 * <p>
050 * This class is thread-safe; the flag is {@code volatile}.
051 *
052 * <h5 class='section'>See Also:</h5><ul>
053 *    <li class='jc'>{@link HealthAggregator}
054 * </ul>
055 *
056 * @since 10.0.0
057 */
058public final class ReadinessState {
059
060   private static final ReadinessState SHARED = new ReadinessState();
061
062   /**
063    * Returns the process-wide shared readiness state.
064    *
065    * <p>
066    * Used as the fallback when no {@code ReadinessState} bean is registered in the relevant {@link BeanStore}.
067    * This is <b>not</b> the standalone Jetty/Tomcat microservice path: that runtime constructs a per-service
068    * instance and publishes it into both the microservice's bean store and the health-probe servlet's bean
069    * store (see the class-level <i>Resolution</i> notes above), so two default-configured microservices in the
070    * same JVM never share readiness.  This shared instance exists only for embeddings that skip that publish
071    * step.
072    *
073    * @return The shared instance.  Never <jk>null</jk>.
074    */
075   public static ReadinessState shared() {
076      return SHARED;
077   }
078
079   /**
080    * Resolves the readiness state from the given bean store, falling back to the {@link #shared() shared} instance.
081    *
082    * @param beanStore The bean store to consult.  Can be <jk>null</jk> (resolves to the shared instance).
083    * @return The resolved readiness state.  Never <jk>null</jk>.
084    */
085   public static ReadinessState resolve(BeanStore beanStore) {
086      if (beanStore == null)
087         return SHARED;
088      return beanStore.getBean(ReadinessState.class).orElse(SHARED);
089   }
090
091   private volatile boolean ready = true;
092   private final ConcurrentMap<Object,Runnable> outOfServiceListeners = new ConcurrentHashMap<>();
093
094   /**
095    * Constructor.
096    *
097    * <p>
098    * The new state starts {@link #isReady() ready}.
099    */
100   public ReadinessState() {
101      // Intentionally empty — the volatile `ready` field defaults to true.
102   }
103
104   /**
105    * Returns whether this service is currently ready to receive traffic.
106    *
107    * @return <jk>true</jk> if ready, <jk>false</jk> if out of service.
108    */
109   public boolean isReady() {
110      return ready;
111   }
112
113   /**
114    * Marks this service as ready to receive traffic.
115    *
116    * <p>
117    * Called when an embedded server starts so a freshly-(re)started service is ready.
118    *
119    * @return This object.
120    */
121   public ReadinessState markReady() {
122      ready = true;
123      return this;
124   }
125
126   /**
127    * Marks this service as out of service so the readiness probe ({@code /readyz}) returns {@code 503}.
128    *
129    * <p>
130    * Called at the very beginning of shutdown, before the listening connector stops.
131    *
132    * @return This object.
133    */
134   public ReadinessState markOutOfService() {
135      ready = false;
136      for (var listener : outOfServiceListeners.values()) {
137         try {
138            listener.run();
139         } catch (RuntimeException e) { // NOSONAR - one failing listener must not stop shutdown or the remaining listeners
140            // Deliberately swallowed: shutdown must proceed regardless of a misbehaving listener.
141         }
142      }
143      return this;
144   }
145
146   /**
147    * Registers a callback to run when this state is {@link #markOutOfService() marked out of service}.
148    *
149    * <p>
150    * Lets components that hold long-lived connections (for example MCP {@code subscriptions/listen} SSE
151    * streams) end them as soon as shutdown begins, instead of making the server's graceful drain wait out its
152    * full stop timeout for streams that would otherwise never finish on their own.
153    *
154    * <p>
155    * Registering a second callback under the same <c>key</c> replaces the first, so a component that is
156    * re-published on every server (re)start does not accumulate duplicate callbacks.  Callbacks run on the
157    * thread that calls {@link #markOutOfService()}; exceptions they throw are swallowed.
158    *
159    * <h5 class='section'>Example:</h5>
160    * <p class='bjava'>
161    *    <jv>state</jv>.onOutOfService(<js>"my-streams"</js>, () -&gt; <jv>streamRegistry</jv>.closeAll());
162    * </p>
163    *
164    * @param key The owner of the callback (typically the registering component).  Must not be <jk>null</jk>.
165    * @param callback The callback to run.  Must not be <jk>null</jk>.
166    * @return This object.
167    */
168   public ReadinessState onOutOfService(Object key, Runnable callback) {
169      outOfServiceListeners.put(key, callback);
170      return this;
171   }
172}