Improve the SPI ClassLoader mechanism to handle more complex scenarios

sentinel-cluster/sentinel-cluster-server-default/src/main/java/com/alibaba/csp/sentinel/cluster/server/TokenServiceProvider.java

@@ -15,13 +15,10 @@
  */
 package com.alibaba.csp.sentinel.cluster.server;
 
-import java.util.ArrayList;
-import java.util.List;
-import java.util.ServiceLoader;
-
 import com.alibaba.csp.sentinel.cluster.TokenService;
 import com.alibaba.csp.sentinel.cluster.flow.DefaultTokenService;
 import com.alibaba.csp.sentinel.log.RecordLog;
+import com.alibaba.csp.sentinel.util.SpiLoader;
 
 /**
  * @author Eric Zhao
@@ -31,8 +28,6 @@ public final class TokenServiceProvider {
 
     private static TokenService service = null;
 
-    private static final ServiceLoader<TokenService> LOADER = ServiceLoader.load(TokenService.class);
-
     static {
         resolveTokenServiceSpi();
     }
@@ -42,24 +37,12 @@ public static TokenService getService() {
     }
 
     private static void resolveTokenServiceSpi() {
-        boolean hasOther = false;
-        List<TokenService> list = new ArrayList<TokenService>();
-        for (TokenService service : LOADER) {
-            if (service.getClass() != DefaultTokenService.class) {
-                hasOther = true;
-                list.add(service);
-            }
-        }
-
-        if (hasOther) {
-            // Pick the first.
-            service = list.get(0);
+        service = SpiLoader.loadFirstInstanceOrDefault(TokenService.class, DefaultTokenService.class);
+        if (service != null) {
+            RecordLog.info("[TokenServiceProvider] Global token service resolved: "
+                + service.getClass().getCanonicalName());
         } else {
-            // No custom token service, using default.
-            service = new DefaultTokenService();
+            RecordLog.warn("[TokenServiceProvider] Unable to resolve TokenService: no SPI found");
         }
-
-        RecordLog.info("[TokenServiceProvider] Global token service resolved: "
-            + service.getClass().getCanonicalName());
     }
 }

sentinel-cluster/sentinel-cluster-server-default/src/main/java/com/alibaba/csp/sentinel/cluster/server/processor/RequestProcessorProvider.java

@@ -20,6 +20,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 
 import com.alibaba.csp.sentinel.cluster.annotation.RequestType;
+import com.alibaba.csp.sentinel.spi.ServiceLoaderUtil;
 import com.alibaba.csp.sentinel.util.AssertUtil;
 
 /**
@@ -30,7 +31,8 @@ public final class RequestProcessorProvider {
 
     private static final Map<Integer, RequestProcessor> PROCESSOR_MAP = new ConcurrentHashMap<>();
 
-    private static final ServiceLoader<RequestProcessor> SERVICE_LOADER = ServiceLoader.load(RequestProcessor.class);
+    private static final ServiceLoader<RequestProcessor> SERVICE_LOADER = ServiceLoaderUtil.getServiceLoader(
+        RequestProcessor.class);
 
     static {
         loadAndInit();
@@ -71,6 +73,5 @@ static void addProcessor(int type, RequestProcessor processor) {
         PROCESSOR_MAP.put(type, processor);
     }
 
-
     private RequestProcessorProvider() {}
 }

sentinel-core/src/main/java/com/alibaba/csp/sentinel/config/SentinelConfig.java

@@ -48,6 +48,7 @@ public class SentinelConfig {
     public static final String TOTAL_METRIC_FILE_COUNT = "csp.sentinel.metric.file.total.count";
     public static final String COLD_FACTOR = "csp.sentinel.flow.cold.factor";
     public static final String STATISTIC_MAX_RT = "csp.sentinel.statistic.max.rt";
+    public static final String SPI_CLASSLOADER = "csp.sentinel.spi.classloader";
 
     static final String DEFAULT_CHARSET = "UTF-8";
     static final long DEFAULT_SINGLE_METRIC_FILE_SIZE = 1024 * 1024 * 50;

sentinel-core/src/main/java/com/alibaba/csp/sentinel/init/InitExecutor.java

@@ -21,6 +21,7 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import com.alibaba.csp.sentinel.log.RecordLog;
+import com.alibaba.csp.sentinel.spi.ServiceLoaderUtil;
 
 /**
  * Load registered init functions and execute in order.
@@ -42,7 +43,7 @@ public static void doInit() {
             return;
         }
         try {
-            ServiceLoader<InitFunc> loader = ServiceLoader.load(InitFunc.class);
+            ServiceLoader<InitFunc> loader = ServiceLoaderUtil.getServiceLoader(InitFunc.class);
             List<OrderWrapper> initList = new ArrayList<OrderWrapper>();
             for (InitFunc initFunc : loader) {
                 RecordLog.info("[InitExecutor] Found init func: " + initFunc.getClass().getCanonicalName());

sentinel-core/src/main/java/com/alibaba/csp/sentinel/slotchain/SlotChainProvider.java

@@ -17,7 +17,7 @@
 
 import com.alibaba.csp.sentinel.log.RecordLog;
 import com.alibaba.csp.sentinel.slots.DefaultSlotChainBuilder;
-import java.util.ServiceLoader;
+import com.alibaba.csp.sentinel.util.SpiLoader;
 
 /**
  * A provider for creating slot chains via resolved slot chain builder SPI.
@@ -29,8 +29,6 @@ public final class SlotChainProvider {
 
     private static volatile SlotChainBuilder slotChainBuilder = null;
 
-    private static final ServiceLoader<SlotChainBuilder> LOADER = ServiceLoader.load(SlotChainBuilder.class);
-
     /**
      * The load and pick process is not thread-safe, but it's okay since the method should be only invoked
      * via {@code lookProcessChain} in {@link com.alibaba.csp.sentinel.CtSph} under lock.
@@ -42,30 +40,19 @@ public static ProcessorSlotChain newSlotChain() {
             return slotChainBuilder.build();
         }
 
-        resolveSlotChainBuilder();
+        // Resolve the slot chain builder SPI.
+        slotChainBuilder = SpiLoader.loadFirstInstanceOrDefault(SlotChainBuilder.class, DefaultSlotChainBuilder.class);
 
         if (slotChainBuilder == null) {
+            // Should not go through here.
             RecordLog.warn("[SlotChainProvider] Wrong state when resolving slot chain builder, using default");
             slotChainBuilder = new DefaultSlotChainBuilder();
+        } else {
+            RecordLog.info("[SlotChainProvider] Global slot chain builder resolved: "
+                + slotChainBuilder.getClass().getCanonicalName());
         }
         return slotChainBuilder.build();
     }
 
-    private static void resolveSlotChainBuilder() {
-        for (SlotChainBuilder builder : LOADER) {
-            if (builder.getClass() != DefaultSlotChainBuilder.class) {
-                slotChainBuilder = builder;
-                break;
-            }
-        }
-        if (slotChainBuilder == null){
-            // No custom builder, using default.
-            slotChainBuilder = new DefaultSlotChainBuilder();
-        }
-
-        RecordLog.info("[SlotChainProvider] Global slot chain builder resolved: "
-            + slotChainBuilder.getClass().getCanonicalName());
-    }
-
     private SlotChainProvider() {}
 }

sentinel-core/src/main/java/com/alibaba/csp/sentinel/spi/ServiceLoaderUtil.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright 1999-2019 Alibaba Group Holding Ltd.
+ *
+ * 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.alibaba.csp.sentinel.spi;
+
+import java.util.ServiceLoader;
+
+import com.alibaba.csp.sentinel.config.SentinelConfig;
+
+/**
+ * @author Eric Zhao
+ * @since 1.7.0
+ */
+public final class ServiceLoaderUtil {
+
+    private static final String CLASSLOADER_DEFAULT = "default";
+    private static final String CLASSLOADER_CONTEXT = "context";
+
+    public static <S> ServiceLoader<S> getServiceLoader(Class<S> clazz) {
+        if (shouldUseContextClassloader()) {
+            return ServiceLoader.load(clazz);
+        } else {
+            return ServiceLoader.load(clazz, clazz.getClassLoader());
+        }
+    }
+
+    public static boolean shouldUseContextClassloader() {
+        String classloaderConf = SentinelConfig.getConfig(SentinelConfig.SPI_CLASSLOADER);
+        return CLASSLOADER_CONTEXT.equalsIgnoreCase(classloaderConf);
+    }
+
+    private ServiceLoaderUtil() {}
+}

sentinel-core/src/main/java/com/alibaba/csp/sentinel/util/SpiLoader.java

@@ -23,6 +23,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 
 import com.alibaba.csp.sentinel.log.RecordLog;
+import com.alibaba.csp.sentinel.spi.ServiceLoaderUtil;
 import com.alibaba.csp.sentinel.spi.SpiOrder;
 
 /**
@@ -34,12 +35,13 @@ public final class SpiLoader {
     private static final Map<String, ServiceLoader> SERVICE_LOADER_MAP = new ConcurrentHashMap<String, ServiceLoader>();
 
     public static <T> T loadFirstInstance(Class<T> clazz) {
+        AssertUtil.notNull(clazz, "SPI class cannot be null");
         try {
             String key = clazz.getName();
             // Not thread-safe, as it's expected to be resolved in a thread-safe context.
             ServiceLoader<T> serviceLoader = SERVICE_LOADER_MAP.get(key);
             if (serviceLoader == null) {
-                serviceLoader = ServiceLoader.load(clazz);
+                serviceLoader = ServiceLoaderUtil.getServiceLoader(clazz);
                 SERVICE_LOADER_MAP.put(key, serviceLoader);
             }
 
@@ -56,6 +58,41 @@ public static <T> T loadFirstInstance(Class<T> clazz) {
         }
     }
 
+    /**
+     * Load the first-found specific SPI instance (excluding provided default SPI class).
+     * If no other SPI implementation found, then create a default SPI instance.
+     *
+     * @param clazz        class of the SPI interface
+     * @param defaultClass class of the default SPI implementation (if no other implementation found)
+     * @param <T>          SPI type
+     * @return the first specific SPI instance if exists, or else the default SPI instance
+     * @since 1.7.0
+     */
+    public static <T> T loadFirstInstanceOrDefault(Class<T> clazz, Class<? extends T> defaultClass) {
+        AssertUtil.notNull(clazz, "SPI class cannot be null");
+        AssertUtil.notNull(defaultClass, "default SPI class cannot be null");
+        try {
+            String key = clazz.getName();
+            // Not thread-safe, as it's expected to be resolved in a thread-safe context.
+            ServiceLoader<T> serviceLoader = SERVICE_LOADER_MAP.get(key);
+            if (serviceLoader == null) {
+                serviceLoader = ServiceLoaderUtil.getServiceLoader(clazz);
+                SERVICE_LOADER_MAP.put(key, serviceLoader);
+            }
+
+            for (T instance : serviceLoader) {
+                if (instance.getClass() != defaultClass) {
+                    return instance;
+                }
+            }
+            return defaultClass.newInstance();
+        } catch (Throwable t) {
+            RecordLog.warn("[SpiLoader] ERROR: loadFirstInstanceOrDefault failed", t);
+            t.printStackTrace();
+            return null;
+        }
+    }
+
     /**
      * Load the SPI instance with highest priority.
      *
@@ -70,7 +107,7 @@ public static <T> T loadHighestPriorityInstance(Class<T> clazz) {
             // Not thread-safe, as it's expected to be resolved in a thread-safe context.
             ServiceLoader<T> serviceLoader = SERVICE_LOADER_MAP.get(key);
             if (serviceLoader == null) {
-                serviceLoader = ServiceLoader.load(clazz);
+                serviceLoader = ServiceLoaderUtil.getServiceLoader(clazz);
                 SERVICE_LOADER_MAP.put(key, serviceLoader);
             }
 
@@ -105,7 +142,7 @@ public static <T> List<T> loadInstanceList(Class<T> clazz) {
             // Not thread-safe, as it's expected to be resolved in a thread-safe context.
             ServiceLoader<T> serviceLoader = SERVICE_LOADER_MAP.get(key);
             if (serviceLoader == null) {
-                serviceLoader = ServiceLoader.load(clazz);
+                serviceLoader = ServiceLoaderUtil.getServiceLoader(clazz);
                 SERVICE_LOADER_MAP.put(key, serviceLoader);
             }
 
@@ -135,7 +172,7 @@ public static <T> List<T> loadInstanceListSorted(Class<T> clazz) {
             // Not thread-safe, as it's expected to be resolved in a thread-safe context.
             ServiceLoader<T> serviceLoader = SERVICE_LOADER_MAP.get(key);
             if (serviceLoader == null) {
-                serviceLoader = ServiceLoader.load(clazz);
+                serviceLoader = ServiceLoaderUtil.getServiceLoader(clazz);
                 SERVICE_LOADER_MAP.put(key, serviceLoader);
             }
 

sentinel-transport/sentinel-transport-common/src/main/java/com/alibaba/csp/sentinel/command/CommandHandlerProvider.java

@@ -21,6 +21,7 @@
 import java.util.ServiceLoader;
 
 import com.alibaba.csp.sentinel.command.annotation.CommandMapping;
+import com.alibaba.csp.sentinel.spi.ServiceLoaderUtil;
 import com.alibaba.csp.sentinel.util.StringUtil;
 
 /**
@@ -30,7 +31,8 @@
  */
 public class CommandHandlerProvider implements Iterable<CommandHandler> {
 
-    private final ServiceLoader<CommandHandler> serviceLoader = ServiceLoader.load(CommandHandler.class);
+    private final ServiceLoader<CommandHandler> serviceLoader = ServiceLoaderUtil.getServiceLoader(
+        CommandHandler.class);
 
     /**
      * Get all command handlers annotated with {@link CommandMapping} with command name.