Allow use SkyWalking plugin to override service in Agent core. (#1111)

* Allow use SkyWalking plugin to override service in Agent core.

* Fix CI.

* Remove duplicate plugin load.

* Add test cases for new plugin boot service machanism.

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/boot/BootService.java

@@ -27,11 +27,11 @@ package org.apache.skywalking.apm.agent.core.boot;
  * @author wusheng
  */
 public interface BootService {
-    void beforeBoot() throws Throwable;
+    void prepare() throws Throwable;
 
     void boot() throws Throwable;
 
-    void afterBoot() throws Throwable;
+    void onComplete() throws Throwable;
 
     void shutdown() throws Throwable;
 }

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/boot/DefaultImplementor.java

+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.skywalking.apm.agent.core.boot;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * @author wusheng
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Target(ElementType.TYPE)
+public @interface DefaultImplementor {
+}

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/boot/OverrideImplementor.java

+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.skywalking.apm.agent.core.boot;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * @author wusheng
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Target(ElementType.TYPE)
+public @interface OverrideImplementor {
+    Class<? extends BootService> value();
+}

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/boot/ServiceConflictException.java

+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.skywalking.apm.agent.core.boot;
+
+/**
+ * @author wusheng
+ */
+public class ServiceConflictException extends RuntimeException {
+    public ServiceConflictException(String message) {
+        super(message);
+    }
+}

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/boot/ServiceManager.java

@@ -16,16 +16,18 @@
  *
  */
 
-
 package org.apache.skywalking.apm.agent.core.boot;
 
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.LinkedHashMap;
+import java.util.LinkedList;
+import java.util.List;
 import java.util.Map;
 import java.util.ServiceLoader;
-import org.apache.skywalking.apm.agent.core.logging.api.LogManager;
 import org.apache.skywalking.apm.agent.core.logging.api.ILog;
+import org.apache.skywalking.apm.agent.core.logging.api.LogManager;
+import org.apache.skywalking.apm.agent.core.plugin.loader.AgentClassLoader;
 
 /**
  * The <code>ServiceManager</code> bases on {@link ServiceLoader},
@@ -42,9 +44,9 @@ public enum ServiceManager {
     public void boot() {
         bootedServices = loadAllServices();
 
-        beforeBoot();
+        prepare();
         startup();
-        afterBoot();
+        onComplete();
     }
 
     public void shutdown() {
@@ -59,18 +61,52 @@ public enum ServiceManager {
 
     private Map<Class, BootService> loadAllServices() {
         Map<Class, BootService> bootedServices = new LinkedHashMap<Class, BootService>();
-        Iterator<BootService> serviceIterator = load().iterator();
+        List<BootService> allServices = new LinkedList<BootService>();
+        load(allServices);
+        Iterator<BootService> serviceIterator = allServices.iterator();
         while (serviceIterator.hasNext()) {
             BootService bootService = serviceIterator.next();
-            bootedServices.put(bootService.getClass(), bootService);
+
+            Class<? extends BootService> bootServiceClass = bootService.getClass();
+            boolean isDefaultImplementor = bootServiceClass.isAnnotationPresent(DefaultImplementor.class);
+            if (isDefaultImplementor) {
+                if (!bootedServices.containsKey(bootServiceClass)) {
+                    bootedServices.put(bootServiceClass, bootService);
+                } else {
+                    //ignore the default service
+                }
+            } else {
+                OverrideImplementor overrideImplementor = bootServiceClass.getAnnotation(OverrideImplementor.class);
+                if (overrideImplementor == null) {
+                    if (!bootedServices.containsKey(bootServiceClass)) {
+                        bootedServices.put(bootServiceClass, bootService);
+                    } else {
+                        throw new ServiceConflictException("Duplicate service define for :" + bootServiceClass);
+                    }
+                } else {
+                    Class<? extends BootService> targetService = overrideImplementor.value();
+                    if (bootedServices.containsKey(targetService)) {
+                        boolean presentDefault = bootedServices.get(targetService).getClass().isAnnotationPresent(DefaultImplementor.class);
+                        if (presentDefault) {
+                            bootedServices.put(targetService, bootService);
+                        } else {
+                            throw new ServiceConflictException("Service " + bootServiceClass + " overrides conflict, " +
+                                "exist more than one service want to override :" + targetService);
+                        }
+                    } else {
+                        bootedServices.put(targetService, bootService);
+                    }
+                }
+            }
+
         }
         return bootedServices;
     }
 
-    private void beforeBoot() {
+    private void prepare() {
         for (BootService service : bootedServices.values()) {
             try {
-                service.beforeBoot();
+                service.prepare();
             } catch (Throwable e) {
                 logger.error(e, "ServiceManager try to pre-start [{}] fail.", service.getClass().getName());
             }
@@ -87,10 +123,10 @@ public enum ServiceManager {
         }
     }
 
-    private void afterBoot() {
+    private void onComplete() {
         for (BootService service : bootedServices.values()) {
             try {
-                service.afterBoot();
+                service.onComplete();
             } catch (Throwable e) {
                 logger.error(e, "Service [{}] AfterBoot process fails.", service.getClass().getName());
             }
@@ -108,7 +144,10 @@ public enum ServiceManager {
         return (T)bootedServices.get(serviceClass);
     }
 
-    ServiceLoader<BootService> load() {
-        return ServiceLoader.load(BootService.class);
+    void load(List<BootService> allServices) {
+        Iterator<BootService> iterator = ServiceLoader.load(BootService.class, AgentClassLoader.getDefault()).iterator();
+        while (iterator.hasNext()) {
+            allServices.add(iterator.next());
+        }
     }
 }

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/ContextManager.java

@@ -20,7 +20,6 @@ package org.apache.skywalking.apm.agent.core.context;
 
 import org.apache.skywalking.apm.agent.core.boot.BootService;
 import org.apache.skywalking.apm.agent.core.boot.ServiceManager;
-import org.apache.skywalking.apm.agent.core.conf.Config;
 import org.apache.skywalking.apm.agent.core.conf.RemoteDownstreamConfig;
 import org.apache.skywalking.apm.agent.core.context.trace.AbstractSpan;
 import org.apache.skywalking.apm.agent.core.context.trace.TraceSegment;
@@ -42,11 +41,14 @@ import org.apache.skywalking.apm.util.StringUtil;
  */
 public class ContextManager implements TracingContextListener, BootService, IgnoreTracerContextListener {
     private static final ILog logger = LogManager.getLogger(ContextManager.class);
-
     private static ThreadLocal<AbstractTracerContext> CONTEXT = new ThreadLocal<AbstractTracerContext>();
+    private static ContextManagerExtendService EXTEND_SERVICE;
 
     private static AbstractTracerContext getOrCreate(String operationName, boolean forceSampling) {
         AbstractTracerContext context = CONTEXT.get();
+        if (EXTEND_SERVICE == null) {
+            EXTEND_SERVICE = ServiceManager.INSTANCE.findService(ContextManagerExtendService.class);
+        }
         if (context == null) {
             if (StringUtil.isEmpty(operationName)) {
                 if (logger.isDebugEnable()) {
@@ -57,17 +59,7 @@ public class ContextManager implements TracingContextListener, BootService, Igno
                 if (RemoteDownstreamConfig.Agent.APPLICATION_ID != DictionaryUtil.nullValue()
                     && RemoteDownstreamConfig.Agent.APPLICATION_INSTANCE_ID != DictionaryUtil.nullValue()
                     ) {
-                    int suffixIdx = operationName.lastIndexOf(".");
-                    if (suffixIdx > -1 && Config.Agent.IGNORE_SUFFIX.contains(operationName.substring(suffixIdx))) {
-                        context = new IgnoredTracerContext();
-                    } else {
-                        SamplingService samplingService = ServiceManager.INSTANCE.findService(SamplingService.class);
-                        if (forceSampling || samplingService.trySampling()) {
-                            context = new TracingContext();
-                        } else {
-                            context = new IgnoredTracerContext();
-                        }
-                    }
+                    context = EXTEND_SERVICE.createTraceContext(operationName, forceSampling);
                 } else {
                     /**
                      * Can't register to collector, no need to trace anything.
@@ -172,18 +164,18 @@ public class ContextManager implements TracingContextListener, BootService, Igno
     }
 
     @Override
-    public void beforeBoot() throws Throwable {
+    public void prepare() throws Throwable {
 
     }
 
     @Override
     public void boot() {
-        TracingContext.ListenerManager.add(this);
-        IgnoredTracerContext.ListenerManager.add(this);
+        ContextManagerExtendService service = ServiceManager.INSTANCE.findService(ContextManagerExtendService.class);
+        service.registerListeners(this);
     }
 
     @Override
-    public void afterBoot() throws Throwable {
+    public void onComplete() throws Throwable {
 
     }
 

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/context/ContextManagerExtendService.java

+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.skywalking.apm.agent.core.context;
+
+import org.apache.skywalking.apm.agent.core.boot.BootService;
+import org.apache.skywalking.apm.agent.core.boot.DefaultImplementor;
+import org.apache.skywalking.apm.agent.core.boot.ServiceManager;
+import org.apache.skywalking.apm.agent.core.conf.Config;
+import org.apache.skywalking.apm.agent.core.sampling.SamplingService;
+
+/**
+ * @author wusheng
+ */
+@DefaultImplementor
+public class ContextManagerExtendService implements BootService {
+    @Override public void prepare() {
+
+    }
+
+    @Override public void boot() {
+
+    }
+
+    @Override public void onComplete() {
+
+    }
+
+    @Override public void shutdown() {
+
+    }
+
+    public void registerListeners(ContextManager manager) {
+        TracingContext.ListenerManager.add(manager);
+        IgnoredTracerContext.ListenerManager.add(manager);
+    }
+
+    public AbstractTracerContext createTraceContext(String operationName, boolean forceSampling) {
+        AbstractTracerContext context;
+        int suffixIdx = operationName.lastIndexOf(".");
+        if (suffixIdx > -1 && Config.Agent.IGNORE_SUFFIX.contains(operationName.substring(suffixIdx))) {
+            context = new IgnoredTracerContext();
+        } else {
+            SamplingService samplingService = ServiceManager.INSTANCE.findService(SamplingService.class);
+            if (forceSampling || samplingService.trySampling()) {
+                context = new TracingContext();
+            } else {
+                context = new IgnoredTracerContext();
+            }
+        }
+
+        return context;
+    }
+}

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/jvm/JVMService.java

@@ -20,6 +20,7 @@ package org.apache.skywalking.apm.agent.core.jvm;
 
 import io.grpc.Channel;
 import org.apache.skywalking.apm.agent.core.boot.BootService;
+import org.apache.skywalking.apm.agent.core.boot.DefaultImplementor;
 import org.apache.skywalking.apm.agent.core.boot.DefaultNamedThreadFactory;
 import org.apache.skywalking.apm.agent.core.boot.ServiceManager;
 import org.apache.skywalking.apm.agent.core.conf.Config;
@@ -52,6 +53,7 @@ import java.util.concurrent.TimeUnit;
  *
  * @author wusheng
  */
+@DefaultImplementor
 public class JVMService implements BootService, Runnable {
     private static final ILog logger = LogManager.getLogger(JVMService.class);
     private LinkedBlockingQueue<JVMMetric> queue;
@@ -60,7 +62,7 @@ public class JVMService implements BootService, Runnable {
     private Sender sender;
 
     @Override
-    public void beforeBoot() throws Throwable {
+    public void prepare() throws Throwable {
         queue = new LinkedBlockingQueue(Config.Jvm.BUFFER_SIZE);
         sender = new Sender();
         ServiceManager.INSTANCE.findService(GRPCChannelManager.class).addChannelListener(sender);
@@ -86,7 +88,7 @@ public class JVMService implements BootService, Runnable {
     }
 
     @Override
-    public void afterBoot() throws Throwable {
+    public void onComplete() throws Throwable {
 
     }
 

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/plugin/loader/AgentClassLoader.java

@@ -16,16 +16,13 @@
  *
  */
 
-
 package org.apache.skywalking.apm.agent.core.plugin.loader;
 
-import org.apache.skywalking.apm.agent.core.boot.AgentPackageNotFoundException;
-import org.apache.skywalking.apm.agent.core.boot.AgentPackagePath;
-import org.apache.skywalking.apm.agent.core.logging.api.LogManager;
-import org.apache.skywalking.apm.agent.core.plugin.PluginBootstrap;
-import org.apache.skywalking.apm.agent.core.logging.api.ILog;
-
-import java.io.*;
+import java.io.BufferedInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.File;
+import java.io.FilenameFilter;
+import java.io.IOException;
 import java.net.MalformedURLException;
 import java.net.URL;
 import java.util.Enumeration;
@@ -35,6 +32,11 @@ import java.util.List;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.jar.JarEntry;
 import java.util.jar.JarFile;
+import org.apache.skywalking.apm.agent.core.boot.AgentPackageNotFoundException;
+import org.apache.skywalking.apm.agent.core.boot.AgentPackagePath;
+import org.apache.skywalking.apm.agent.core.logging.api.ILog;
+import org.apache.skywalking.apm.agent.core.logging.api.LogManager;
+import org.apache.skywalking.apm.agent.core.plugin.PluginBootstrap;
 
 /**
  * The <code>AgentClassLoader</code> represents a classloader,
@@ -64,7 +66,13 @@ public class AgentClassLoader extends ClassLoader {
      * @throws AgentPackageNotFoundException
      */
     public static AgentClassLoader initDefaultLoader() throws AgentPackageNotFoundException {
-        DEFAULT_LOADER = new AgentClassLoader(PluginBootstrap.class.getClassLoader());
+        if (DEFAULT_LOADER == null) {
+            synchronized (AgentClassLoader.class) {
+                if (DEFAULT_LOADER == null) {
+                    DEFAULT_LOADER = new AgentClassLoader(PluginBootstrap.class.getClassLoader());
+                }
+            }
+        }
         return getDefault();
     }
 

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/remote/AppAndServiceRegisterClient.java

@@ -20,6 +20,7 @@ package org.apache.skywalking.apm.agent.core.remote;
 
 import io.grpc.Channel;
 import org.apache.skywalking.apm.agent.core.boot.BootService;
+import org.apache.skywalking.apm.agent.core.boot.DefaultImplementor;
 import org.apache.skywalking.apm.agent.core.boot.DefaultNamedThreadFactory;
 import org.apache.skywalking.apm.agent.core.boot.ServiceManager;
 import org.apache.skywalking.apm.agent.core.conf.Config;
@@ -44,6 +45,7 @@ import java.util.concurrent.TimeUnit;
 /**
  * @author wusheng
  */
+@DefaultImplementor
 public class AppAndServiceRegisterClient implements BootService, GRPCChannelListener, Runnable, TracingContextListener {
     private static final ILog logger = LogManager.getLogger(AppAndServiceRegisterClient.class);
     private static final String PROCESS_UUID = UUID.randomUUID().toString().replaceAll("-", "");
@@ -73,7 +75,7 @@ public class AppAndServiceRegisterClient implements BootService, GRPCChannelList
     }
 
     @Override
-    public void beforeBoot() throws Throwable {
+    public void prepare() throws Throwable {
         ServiceManager.INSTANCE.findService(GRPCChannelManager.class).addChannelListener(this);
     }
 
@@ -89,7 +91,7 @@ public class AppAndServiceRegisterClient implements BootService, GRPCChannelList
     }
 
     @Override
-    public void afterBoot() throws Throwable {
+    public void onComplete() throws Throwable {
         TracingContext.ListenerManager.add(this);
     }
 

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/remote/CollectorDiscoveryService.java

@@ -19,6 +19,7 @@
 package org.apache.skywalking.apm.agent.core.remote;
 
 import org.apache.skywalking.apm.agent.core.boot.BootService;
+import org.apache.skywalking.apm.agent.core.boot.DefaultImplementor;
 import org.apache.skywalking.apm.agent.core.boot.DefaultNamedThreadFactory;
 import org.apache.skywalking.apm.agent.core.conf.Config;
 import org.apache.skywalking.apm.agent.core.conf.RemoteDownstreamConfig;
@@ -36,12 +37,13 @@ import java.util.concurrent.TimeUnit;
  *
  * @author wusheng
  */
+@DefaultImplementor
 public class CollectorDiscoveryService implements BootService {
     private static final ILog logger = LogManager.getLogger(CollectorDiscoveryService.class);
     private ScheduledFuture<?> future;
 
     @Override
-    public void beforeBoot() {
+    public void prepare() {
 
     }
 
@@ -72,7 +74,7 @@ public class CollectorDiscoveryService implements BootService {
     }
 
     @Override
-    public void afterBoot() throws Throwable {
+    public void onComplete() throws Throwable {
 
     }
 

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/remote/GRPCChannelManager.java

@@ -29,6 +29,7 @@ import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 import org.apache.skywalking.apm.agent.core.boot.BootService;
+import org.apache.skywalking.apm.agent.core.boot.DefaultImplementor;
 import org.apache.skywalking.apm.agent.core.boot.DefaultNamedThreadFactory;
 import org.apache.skywalking.apm.agent.core.conf.Config;
 import org.apache.skywalking.apm.agent.core.conf.RemoteDownstreamConfig;
@@ -39,6 +40,7 @@ import org.apache.skywalking.apm.util.RunnableWithExceptionProtection;
 /**
  * @author wusheng, zhang xin
  */
+@DefaultImplementor
 public class GRPCChannelManager implements BootService, Runnable {
     private static final ILog logger = LogManager.getLogger(GRPCChannelManager.class);
 
@@ -49,7 +51,7 @@ public class GRPCChannelManager implements BootService, Runnable {
     private List<GRPCChannelListener> listeners = Collections.synchronizedList(new LinkedList<GRPCChannelListener>());
 
     @Override
-    public void beforeBoot() throws Throwable {
+    public void prepare() throws Throwable {
 
     }
 
@@ -66,7 +68,7 @@ public class GRPCChannelManager implements BootService, Runnable {
     }
 
     @Override
-    public void afterBoot() throws Throwable {
+    public void onComplete() throws Throwable {
 
     }
 

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/remote/TraceSegmentServiceClient.java

@@ -22,6 +22,7 @@ package org.apache.skywalking.apm.agent.core.remote;
 import io.grpc.Channel;
 import io.grpc.stub.StreamObserver;
 import org.apache.skywalking.apm.agent.core.boot.BootService;
+import org.apache.skywalking.apm.agent.core.boot.DefaultImplementor;
 import org.apache.skywalking.apm.agent.core.boot.ServiceManager;
 import org.apache.skywalking.apm.agent.core.context.TracingContext;
 import org.apache.skywalking.apm.agent.core.context.TracingContextListener;
@@ -44,6 +45,7 @@ import static org.apache.skywalking.apm.agent.core.remote.GRPCChannelStatus.CONN
 /**
  * @author wusheng
  */
+@DefaultImplementor
 public class TraceSegmentServiceClient implements BootService, IConsumer<TraceSegment>, TracingContextListener, GRPCChannelListener {
     private static final ILog logger = LogManager.getLogger(TraceSegmentServiceClient.class);
     private static final int TIMEOUT = 30 * 1000;
@@ -56,7 +58,7 @@ public class TraceSegmentServiceClient implements BootService, IConsumer<TraceSe
     private volatile GRPCChannelStatus status = GRPCChannelStatus.DISCONNECT;
 
     @Override
-    public void beforeBoot() throws Throwable {
+    public void prepare() throws Throwable {
         ServiceManager.INSTANCE.findService(GRPCChannelManager.class).addChannelListener(this);
     }
 
@@ -71,7 +73,7 @@ public class TraceSegmentServiceClient implements BootService, IConsumer<TraceSe
     }
 
     @Override
-    public void afterBoot() throws Throwable {
+    public void onComplete() throws Throwable {
         TracingContext.ListenerManager.add(this);
     }
 

apm-sniffer/apm-agent-core/src/main/java/org/apache/skywalking/apm/agent/core/sampling/SamplingService.java

@@ -25,6 +25,7 @@ import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import org.apache.skywalking.apm.agent.core.boot.BootService;
+import org.apache.skywalking.apm.agent.core.boot.DefaultImplementor;
 import org.apache.skywalking.apm.agent.core.boot.DefaultNamedThreadFactory;
 import org.apache.skywalking.apm.agent.core.conf.Config;
 import org.apache.skywalking.apm.agent.core.context.trace.TraceSegment;
@@ -41,6 +42,7 @@ import org.apache.skywalking.apm.util.RunnableWithExceptionProtection;
  *
  * @author wusheng
  */
+@DefaultImplementor
 public class SamplingService implements BootService {
     private static final ILog logger = LogManager.getLogger(SamplingService.class);
 
@@ -49,7 +51,7 @@ public class SamplingService implements BootService {
     private volatile ScheduledFuture<?> scheduledFuture;
 
     @Override
-    public void beforeBoot() throws Throwable {
+    public void prepare() throws Throwable {
 
     }
 
@@ -82,7 +84,7 @@ public class SamplingService implements BootService {
     }
 
     @Override
-    public void afterBoot() throws Throwable {
+    public void onComplete() throws Throwable {
 
     }
 

apm-sniffer/apm-agent-core/src/main/resources/META-INF/services/org.apache.skywalking.apm.agent.core.boot.BootService

@@ -23,3 +23,4 @@ org.apache.skywalking.apm.agent.core.sampling.SamplingService
 org.apache.skywalking.apm.agent.core.remote.GRPCChannelManager
 org.apache.skywalking.apm.agent.core.jvm.JVMService
 org.apache.skywalking.apm.agent.core.remote.AppAndServiceRegisterClient
+org.apache.skywalking.apm.agent.core.context.ContextManagerExtendService

apm-sniffer/apm-agent-core/src/test/java/org/apache/skywalking/apm/agent/core/boot/ServiceManagerTest.java

@@ -57,7 +57,7 @@ public class ServiceManagerTest {
     public void testServiceDependencies() throws Exception {
         HashMap<Class, BootService> registryService = getFieldValue(ServiceManager.INSTANCE, "bootedServices");
 
-        assertThat(registryService.size(), is(7));
+        assertThat(registryService.size(), is(8));
 
         assertTraceSegmentServiceClient(ServiceManager.INSTANCE.findService(TraceSegmentServiceClient.class));
         assertContextManager(ServiceManager.INSTANCE.findService(ContextManager.class));

apm-sniffer/apm-sdk-plugin/dubbo-plugin/src/test/java/org/apache/skywalking/apm/plugin/dubbo/ContextManagerExtendOverrideService.java

+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.skywalking.apm.plugin.dubbo;
+
+import org.apache.skywalking.apm.agent.core.boot.OverrideImplementor;
+import org.apache.skywalking.apm.agent.core.context.ContextManagerExtendService;
+
+/**
+ * @author wusheng
+ */
+@OverrideImplementor(ContextManagerExtendService.class)
+public class ContextManagerExtendOverrideService extends ContextManagerExtendService {
+}

apm-sniffer/apm-sdk-plugin/dubbo-plugin/src/test/java/org/apache/skywalking/apm/plugin/dubbo/DubboInterceptorTest.java

@@ -16,7 +16,6 @@
  *
  */
 
-
 package org.apache.skywalking.apm.plugin.dubbo;
 
 import com.alibaba.dubbo.common.URL;
@@ -25,7 +24,9 @@ import com.alibaba.dubbo.rpc.Invoker;
 import com.alibaba.dubbo.rpc.Result;
 import com.alibaba.dubbo.rpc.RpcContext;
 import java.util.List;
+import org.apache.skywalking.apm.agent.core.boot.ServiceManager;
 import org.apache.skywalking.apm.agent.core.conf.Config;
+import org.apache.skywalking.apm.agent.core.context.ContextManagerExtendService;
 import org.apache.skywalking.apm.agent.core.context.SW3CarrierItem;
 import org.apache.skywalking.apm.agent.core.context.trace.AbstractTracingSpan;
 import org.apache.skywalking.apm.agent.core.context.trace.LogDataEntity;
@@ -52,6 +53,7 @@ import org.powermock.api.mockito.PowerMockito;
 import org.powermock.core.classloader.annotations.PrepareForTest;
 import org.powermock.modules.junit4.PowerMockRunner;
 import org.powermock.modules.junit4.PowerMockRunnerDelegate;
+import org.springframework.util.Assert;
 
 import static org.hamcrest.CoreMatchers.is;
 import static org.junit.Assert.assertNull;
@@ -106,6 +108,20 @@ public class DubboInterceptorTest {
         Config.Agent.APPLICATION_CODE = "DubboTestCases-APP";
     }
 
+    @Test
+    public void testServiceFromPlugin() {
+        PluginBootService service = ServiceManager.INSTANCE.findService(PluginBootService.class);
+
+        Assert.notNull(service);
+    }
+
+    @Test
+    public void testServiceOverrideFromPlugin() {
+        ContextManagerExtendService service = ServiceManager.INSTANCE.findService(ContextManagerExtendService.class);
+
+        Assert.isInstanceOf(ContextManagerExtendOverrideService.class, service);
+    }
+
     @Test
     public void testConsumerWithAttachment() throws Throwable {
         dubboInterceptor.beforeMethod(enhancedInstance, null, allArguments, argumentTypes, methodInterceptResult);

apm-sniffer/apm-sdk-plugin/dubbo-plugin/src/test/java/org/apache/skywalking/apm/plugin/dubbo/PluginBootService.java

+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.skywalking.apm.plugin.dubbo;
+
+import org.apache.skywalking.apm.agent.core.boot.BootService;
+
+/**
+ * @author wusheng
+ */
+public class PluginBootService implements BootService {
+    @Override public void prepare() throws Throwable {
+
+    }
+
+    @Override public void boot() throws Throwable {
+
+    }
+
+    @Override public void onComplete() throws Throwable {
+
+    }
+
+    @Override public void shutdown() throws Throwable {
+
+    }
+}

apm-sniffer/apm-sdk-plugin/dubbo-plugin/src/test/resources/META-INF/services/org.apache.skywalking.apm.agent.core.boot.BootService

+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You 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.
+#
+#
+
+org.apache.skywalking.apm.plugin.dubbo.PluginBootService
+org.apache.skywalking.apm.plugin.dubbo.ContextManagerExtendOverrideService

apm-sniffer/apm-test-tools/src/main/java/org/apache/skywalking/apm/agent/test/tools/AgentServiceRule.java

@@ -23,6 +23,7 @@ import java.util.HashMap;
 import java.util.LinkedList;
 import org.apache.skywalking.apm.agent.core.conf.RemoteDownstreamConfig;
 import org.apache.skywalking.apm.agent.core.context.TracingContext;
+import org.apache.skywalking.apm.agent.core.plugin.loader.AgentClassLoader;
 import org.junit.rules.ExternalResource;
 import org.apache.skywalking.apm.agent.core.boot.BootService;
 import org.apache.skywalking.apm.agent.core.boot.ServiceManager;
@@ -49,6 +50,7 @@ public class AgentServiceRule extends ExternalResource {
     @Override
     protected void before() throws Throwable {
         super.before();
+        AgentClassLoader.initDefaultLoader();
         Config.Logging.LEVEL = LogLevel.OFF;
         ServiceManager.INSTANCE.boot();
         RemoteDownstreamConfig.Agent.APPLICATION_ID = 1;