Expose CachingStream as a public API

xds/src/main/java/com/linecorp/armeria/xds/RouteStream.java

@@ -21,6 +21,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.function.Function;
 
 import com.google.common.collect.ImmutableList;
 import com.google.protobuf.Any;
@@ -86,19 +87,19 @@ protected Subscription onStart(SnapshotWatcher<RouteSnapshot> watcher) {
     }
 
     private static final class FilterCaches {
-        final CachingStream<Map<String, Any>, ClientPreprocessors> downstream;
-        final CachingStream<Map<String, Any>, ClientDecoration> upstream;
-        final CachingStream<Map<String, Any>, Optional<HttpService>> server;
+        final Function<Map<String, Any>, SnapshotStream<ClientPreprocessors>> downstream;
+        final Function<Map<String, Any>, SnapshotStream<ClientDecoration>> upstream;
+        final Function<Map<String, Any>, SnapshotStream<Optional<HttpService>>> server;
 
         FilterCaches(XdsExtensionRegistry registry, SubscriptionContext context,
                      List<HttpFilter> hcmHttpFilters, List<HttpFilter> upstreamFilters) {
-            downstream = new CachingStream<>(
+            downstream = SnapshotStream.caching(
                     filterConfigs -> FilterUtil.buildDownstreamFilter(
                             registry, context, hcmHttpFilters, filterConfigs));
-            upstream = new CachingStream<>(
+            upstream = SnapshotStream.caching(
                     filterConfigs -> FilterUtil.buildUpstreamFilter(
                             registry, context, upstreamFilters, filterConfigs));
-            server = new CachingStream<>(
+            server = SnapshotStream.caching(
                     filterConfigs -> FilterUtil.buildDownstreamServerFilter(
                             registry, context, hcmHttpFilters, filterConfigs));
         }
@@ -252,12 +253,12 @@ protected Subscription onStart(SnapshotWatcher<RouteEntry> watcher) {
 
             // Build filter streams (deduped via caches for routes with identical configs)
             final SnapshotStream<ClientPreprocessors> downstreamStream =
-                    filterCaches.downstream.subscribe(filterConfigs);
+                    filterCaches.downstream.apply(filterConfigs);
             final SnapshotStream<ClientDecoration> upstreamStream =
-                    filterCaches.upstream.subscribe(filterConfigs);
+                    filterCaches.upstream.apply(filterConfigs);
 
             final SnapshotStream<Optional<HttpService>> httpServiceStream =
-                    filterCaches.server.subscribe(filterConfigs);
+                    filterCaches.server.apply(filterConfigs);
 
             if (!route.getRoute().hasCluster()) {
                 return SnapshotStream.combineLatest(

xds/src/main/java/com/linecorp/armeria/xds/CachingStream.java → xds/src/main/java/com/linecorp/armeria/xds/stream/CachingStream.java

@@ -14,31 +14,29 @@
  * under the License.
  */
 
-package com.linecorp.armeria.xds;
+package com.linecorp.armeria.xds.stream;
 
 import static java.util.Objects.requireNonNull;
 
 import java.util.HashMap;
 import java.util.Map;
 import java.util.function.Function;
 
-import com.linecorp.armeria.xds.stream.RefCountedStream;
-import com.linecorp.armeria.xds.stream.SnapshotStream;
-import com.linecorp.armeria.xds.stream.Subscription;
+import com.linecorp.armeria.xds.SnapshotWatcher;
 
 final class CachingStream<K, T> {
 
-    private final Function<K, SnapshotStream<T>> factory;
+    private final Function<? super K, ? extends SnapshotStream<T>> factory;
     private final Map<K, CacheEntry> cache = new HashMap<>();
 
-    CachingStream(Function<K, SnapshotStream<T>> factory) {
+    CachingStream(Function<? super K, ? extends SnapshotStream<T>> factory) {
         this.factory = requireNonNull(factory, "factory");
     }
 
     SnapshotStream<T> subscribe(K key) {
         requireNonNull(key, "key");
         return watcher -> {
-            final CacheEntry entry = cache.computeIfAbsent(key, k -> new CacheEntry(k));
+            final CacheEntry entry = cache.computeIfAbsent(key, CacheEntry::new);
             return entry.subscribe(watcher);
         };
     }

xds/src/main/java/com/linecorp/armeria/xds/stream/SnapshotStream.java

@@ -213,6 +213,34 @@ static <T> SnapshotStream<T> error(Throwable error) {
         return new StaticSnapshotStream<>(null, error);
     }
 
+    /**
+     * Returns a caching function that deduplicates {@link SnapshotStream} subscriptions by key
+     * using reference counting. When multiple subscribers request the same key, they share
+     * a single upstream {@link SnapshotStream}. The upstream is created lazily on the first
+     * subscription and closed automatically when the last subscriber unsubscribes.
+     *
+     * <p>Example usage:
+     * <pre>{@code
+     * Function<String, SnapshotStream<MySnapshot>> cached = SnapshotStream.caching(
+     *     name -> createSnapshotStream(name));
+     *
+     * // Both streams share the same underlying subscription for "foo"
+     * SnapshotStream<MySnapshot> stream1 = cached.apply("foo");
+     * SnapshotStream<MySnapshot> stream2 = cached.apply("foo");
+     * }</pre>
+     *
+     * @param factory a function that creates a new {@link SnapshotStream} for a given key
+     * @param <K> the key type used to identify cached streams
+     * @param <T> the type of snapshot values delivered by the cached streams
+     * @return a function that returns cached {@link SnapshotStream}s by key
+     */
+    static <K, T> Function<K, SnapshotStream<T>> caching(
+            Function<? super K, ? extends SnapshotStream<T>> factory) {
+        requireNonNull(factory, "factory");
+        final CachingStream<K, T> cachingStream = new CachingStream<>(factory);
+        return cachingStream::subscribe;
+    }
+
     /**
      * Returns a new stream that asserts {@link #subscribe} and {@link Subscription#close()}
      * are called from the given event loop. Throws {@link IllegalStateException} if called