From 24e731ac15e3ca2f6b0fc645036eec0307c2b195 Mon Sep 17 00:00:00 2001
From: ZHANG Dapeng <zdapeng@google.com>
Date: Thu, 21 Nov 2019 12:37:48 -0800
Subject: [PATCH] remove getLoadStatsStore method from LocalityStore

Currently `LoadStatsStore` is create by `LocalityStore` and `LoadReportClient` retrieves `LoadStatsStore` from `LocalityStore.getLoadStatsStore()`.

But `LocalityStore` is create by EDS policy, whereas `LoadReportClient` and `LoadStatsStore` should be created by CDS (if not EDS-only), before `LocalityStore` is created. If `LoadReportClient` is embedded in `XdsClientImpl`, it need a `LoadStatsStore` which shouldn't be created by `LocalityStore`.

Instead, `LoadStatsStore` should be create before `LocalityStore` is created, and be passed to `LocalityStore`'s constructor. A getter is not needed.
---
 .../main/java/io/grpc/xds/LocalityStore.java  | 20 ++++++++-----------
 .../main/java/io/grpc/xds/LookasideLb.java    |  6 ++++--
 .../java/io/grpc/xds/LookasideLbTest.java     |  5 ++---
 3 files changed, 14 insertions(+), 17 deletions(-)

diff --git a/xds/src/main/java/io/grpc/xds/LocalityStore.java b/xds/src/main/java/io/grpc/xds/LocalityStore.java
index 8f9d396fcc1b..2527c0dc2cee 100644
--- a/xds/src/main/java/io/grpc/xds/LocalityStore.java
+++ b/xds/src/main/java/io/grpc/xds/LocalityStore.java
@@ -76,15 +76,14 @@ interface LocalityStore {
 
   void updateOobMetricsReportInterval(long reportIntervalNano);
 
-  LoadStatsStore getLoadStatsStore();
-
   @VisibleForTesting
   abstract class LocalityStoreFactory {
     private static final LocalityStoreFactory DEFAULT_INSTANCE =
         new LocalityStoreFactory() {
           @Override
-          LocalityStore newLocalityStore(Helper helper, LoadBalancerRegistry lbRegistry) {
-            return new LocalityStoreImpl(helper, lbRegistry);
+          LocalityStore newLocalityStore(
+              Helper helper, LoadBalancerRegistry lbRegistry, LoadStatsStore loadStatsStore) {
+            return new LocalityStoreImpl(helper, lbRegistry, loadStatsStore);
           }
         };
 
@@ -92,7 +91,8 @@ static LocalityStoreFactory getInstance() {
       return DEFAULT_INSTANCE;
     }
 
-    abstract LocalityStore newLocalityStore(Helper helper, LoadBalancerRegistry lbRegistry);
+    abstract LocalityStore newLocalityStore(
+        Helper helper, LoadBalancerRegistry lbRegistry, LoadStatsStore loadStatsStore);
   }
 
   final class LocalityStoreImpl implements LocalityStore {
@@ -113,9 +113,10 @@ final class LocalityStoreImpl implements LocalityStore {
     private List<DropOverload> dropOverloads = ImmutableList.of();
     private long metricsReportIntervalNano = -1;
 
-    LocalityStoreImpl(Helper helper, LoadBalancerRegistry lbRegistry) {
+    LocalityStoreImpl(
+        Helper helper, LoadBalancerRegistry lbRegistry, LoadStatsStore loadStatsStore) {
       this(helper, pickerFactoryImpl, lbRegistry, ThreadSafeRandom.ThreadSafeRandomImpl.instance,
-          new LoadStatsStoreImpl(), OrcaPerRequestUtil.getInstance(), OrcaOobUtil.getInstance());
+          loadStatsStore, OrcaPerRequestUtil.getInstance(), OrcaOobUtil.getInstance());
     }
 
     @VisibleForTesting
@@ -284,11 +285,6 @@ public String toString() {
           TimeUnit.MINUTES, helper.getScheduledExecutorService());
     }
 
-    @Override
-    public LoadStatsStore getLoadStatsStore() {
-      return loadStatsStore;
-    }
-
     @Override
     public void updateOobMetricsReportInterval(long reportIntervalNano) {
       metricsReportIntervalNano = reportIntervalNano;
diff --git a/xds/src/main/java/io/grpc/xds/LookasideLb.java b/xds/src/main/java/io/grpc/xds/LookasideLb.java
index 9fd9a8c8cf3e..d32530c64238 100644
--- a/xds/src/main/java/io/grpc/xds/LookasideLb.java
+++ b/xds/src/main/java/io/grpc/xds/LookasideLb.java
@@ -175,11 +175,13 @@ public void handleResolvedAddresses(ResolvedAddresses resolvedAddresses) {
                 xdsClient  = new XdsComms2(
                     channel, helper, new ExponentialBackoffPolicy.Provider(),
                     GrpcUtil.STOPWATCH_SUPPLIER, node);
-                localityStore = localityStoreFactory.newLocalityStore(helper, lbRegistry);
+                LoadStatsStore loadStatsStore = new LoadStatsStoreImpl();
+                localityStore = localityStoreFactory.newLocalityStore(
+                    helper, lbRegistry, loadStatsStore);
                 // TODO(zdapeng): Use XdsClient to do Lrs directly.
                 lrsClient = loadReportClientFactory.createLoadReportClient(
                     channel, helper, new ExponentialBackoffPolicy.Provider(),
-                    localityStore.getLoadStatsStore());
+                    loadStatsStore);
                 final LoadReportCallback lrsCallback =
                     new LoadReportCallback() {
                       @Override
diff --git a/xds/src/test/java/io/grpc/xds/LookasideLbTest.java b/xds/src/test/java/io/grpc/xds/LookasideLbTest.java
index d32b1dcb8e7c..0f8ff7627028 100644
--- a/xds/src/test/java/io/grpc/xds/LookasideLbTest.java
+++ b/xds/src/test/java/io/grpc/xds/LookasideLbTest.java
@@ -194,11 +194,10 @@ public void onCompleted() {
 
     localityStoreFactory = new LocalityStoreFactory() {
       @Override
-      public LocalityStore newLocalityStore(Helper helper, LoadBalancerRegistry lbRegistry) {
+      public LocalityStore newLocalityStore(
+          Helper helper, LoadBalancerRegistry lbRegistry, LoadStatsStore loadStatsStore) {
         helpers.add(helper);
         LocalityStore localityStore = mock(LocalityStore.class);
-        LoadStatsStore loadStatsStore = mock(LoadStatsStore.class);
-        doReturn(loadStatsStore).when(localityStore).getLoadStatsStore();
         localityStores.add(localityStore);
         return localityStore;
       }