code: /*
   * 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.oap.server.analyzer.provider.trace;
  
  import org.apache.skywalking.oap.server.analyzer.provider.AnalyzerModuleConfig;
  import org.apache.skywalking.oap.server.analyzer.provider.AnalyzerModuleProvider;
  import org.apache.skywalking.oap.server.analyzer.provider.trace.sampling.SamplingPolicy;
  import org.apache.skywalking.oap.server.analyzer.provider.trace.sampling.SamplingPolicySettings;
  import org.apache.skywalking.oap.server.configuration.api.ConfigChangeWatcher;
  import org.apache.skywalking.oap.server.configuration.api.ConfigTable;
  import org.apache.skywalking.oap.server.configuration.api.ConfigWatcherRegister;
  import org.apache.skywalking.oap.server.configuration.api.GroupConfigTable;
  import org.junit.jupiter.api.Assertions;
  import org.junit.jupiter.api.BeforeEach;
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.api.Timeout;
  import org.junit.jupiter.api.extension.ExtendWith;
  import org.mockito.junit.jupiter.MockitoExtension;
  import org.powermock.reflect.Whitebox;
  
  import java.util.Optional;
  import java.util.Set;
  import java.util.concurrent.atomic.AtomicReference;
  
  @ExtendWith(MockitoExtension.class)
  public class TraceSamplingPolicyWatcherTest {
  
      private AnalyzerModuleProvider provider;
      private AnalyzerModuleConfig moduleConfig;
  
      @BeforeEach
      public void init() {
          provider = new AnalyzerModuleProvider();
          moduleConfig = new AnalyzerModuleConfig();
          moduleConfig.setTraceSamplingPolicySettingsFile("trace-sampling-policy-settings.yml");
      }
  
      @Test
      public void testStaticConfigInit() {
          TraceSamplingPolicyWatcher watcher = new TraceSamplingPolicyWatcher(moduleConfig, provider);
          // default sample = 10000
          globalDefaultSamplingRateEquals(watcher, 9999);
      }
  
      @Test
      @Timeout(20)
      public void testTraceLatencyThresholdDynamicUpdate() throws InterruptedException {
          ConfigWatcherRegister register = new TraceLatencyThresholdMockConfigWatcherRegister(3);
  
          TraceSamplingPolicyWatcher watcher = new TraceSamplingPolicyWatcher(moduleConfig, provider);
          register.registerConfigChangeWatcher(watcher);
          register.start();
          // Default duration is -1, so 3000 must not be sampled,until updating to 3000
          while (!watcher.shouldSample("", 10000, 3000)) {
              Thread.sleep(2000);
          }
          Assertions.assertTrue(watcher.shouldSample("", 10000, 3001));
      }
  
      @Test
      public void testTraceLatencyThresholdNotify() {
          TraceSamplingPolicyWatcher watcher = new TraceSamplingPolicyWatcher(moduleConfig, provider);
          ConfigChangeWatcher.ConfigChangeEvent value1 = new ConfigChangeWatcher.ConfigChangeEvent(
              "default:\n" +
                  "  duration: 8000", ConfigChangeWatcher.EventType.MODIFY);
  
          watcher.notify(value1);
          globalDefaultDurationEquals(watcher, 8000);
          Assertions.assertEquals(watcher.value(), "default:\n" +
              "  duration: 8000");
  
          ConfigChangeWatcher.ConfigChangeEvent value2 = new ConfigChangeWatcher.ConfigChangeEvent(
              "default:\n" +
                  "  duration: 8000", ConfigChangeWatcher.EventType.DELETE);
  
          watcher.notify(value2);
          Assertions.assertEquals(watcher.value(), null);
  
          ConfigChangeWatcher.ConfigChangeEvent value3 = new ConfigChangeWatcher.ConfigChangeEvent(
              "default:\n" +
                  "  duration: 800", ConfigChangeWatcher.EventType.ADD);
  
          watcher.notify(value3);
          globalDefaultDurationEquals(watcher, 800);
          Assertions.assertEquals(watcher.value(), "default:\n" +
              "  duration: 800");
  
          ConfigChangeWatcher.ConfigChangeEvent value4 = new ConfigChangeWatcher.ConfigChangeEvent(
              "default:\n" +
                  "  duration: abc", ConfigChangeWatcher.EventType.MODIFY);
  
          watcher.notify(value4);
          globalDefaultDurationEquals(watcher, 800);
          Assertions.assertEquals(watcher.value(), "default:\n" +
              "  duration: 800");
  
          ConfigChangeWatcher.ConfigChangeEvent value5 = new ConfigChangeWatcher.ConfigChangeEvent(
              "default:\n" +
                  "  rate: abc\n" +
                  "  duration: 900", ConfigChangeWatcher.EventType.MODIFY);
  
          watcher.notify(value5);
          globalDefaultDurationEquals(watcher, 800);
          Assertions.assertEquals(watcher.value(), "default:\n" +
              "  duration: 800");
      }
  
      public static class TraceLatencyThresholdMockConfigWatcherRegister extends ConfigWatcherRegister {
  
          public TraceLatencyThresholdMockConfigWatcherRegister(long syncPeriod) {
              super(syncPeriod);
          }
  
          @Override
          public Optional<ConfigTable> readConfig(Set<String> keys) {
              ConfigTable table = new ConfigTable();
              table.add(new ConfigTable.ConfigItem("agent-analyzer.default.traceSamplingPolicy", "default:\n" +
                  "  duration: 3000"));
              return Optional.of(table);
          }
  
          @Override
          public Optional<GroupConfigTable> readGroupConfig(final Set<String> keys) {
              return Optional.empty();
          }
      }
  
      @Test
      @Timeout(20)
      public void testDefaultSampleRateDynamicUpdate() throws InterruptedException {
          ConfigWatcherRegister register = new DefaultSampleRateMockConfigWatcherRegister(3);
  
          TraceSamplingPolicyWatcher watcher = new TraceSamplingPolicyWatcher(moduleConfig, provider);
          register.registerConfigChangeWatcher(watcher);
          register.start();
          // Default is 10000, so 9000 must be sampled,until updating to 9000
          while (watcher.shouldSample("", 9000, -1)) {
              Thread.sleep(2000);
          }
          globalDefaultSamplingRateEquals(watcher, 8999);
      }
  
      @Test
      public void testDefaultSampleRateNotify() {
          TraceSamplingPolicyWatcher watcher = new TraceSamplingPolicyWatcher(moduleConfig, provider);
          ConfigChangeWatcher.ConfigChangeEvent value1 = new ConfigChangeWatcher.ConfigChangeEvent(
              "default:\n" +
                  "  rate: 8000", ConfigChangeWatcher.EventType.MODIFY);
  
          watcher.notify(value1);
          globalDefaultSamplingRateEquals(watcher, 7999);
          Assertions.assertEquals(watcher.value(), "default:\n" +
              "  rate: 8000");
  
          ConfigChangeWatcher.ConfigChangeEvent value2 = new ConfigChangeWatcher.ConfigChangeEvent(
              "default:\n" +
                  "  rate: 1000", ConfigChangeWatcher.EventType.DELETE);
  
          watcher.notify(value2);
          globalDefaultSamplingRateEquals(watcher, 9999);
          Assertions.assertEquals(watcher.value(), null);
  
          ConfigChangeWatcher.ConfigChangeEvent value3 = new ConfigChangeWatcher.ConfigChangeEvent(
              "default:\n" +
                  "  rate: 500", ConfigChangeWatcher.EventType.ADD);
  
          watcher.notify(value3);
          globalDefaultSamplingRateEquals(watcher, 499);
          Assertions.assertEquals(watcher.value(), "default:\n" +
              "  rate: 500");
  
          ConfigChangeWatcher.ConfigChangeEvent value4 = new ConfigChangeWatcher.ConfigChangeEvent(
              "default:\n" +
                  "  rate: abc", ConfigChangeWatcher.EventType.MODIFY);
  
          watcher.notify(value4);
          globalDefaultSamplingRateEquals(watcher, 499);
          Assertions.assertEquals(watcher.value(), "default:\n" +
              "  rate: 500");
  
          ConfigChangeWatcher.ConfigChangeEvent value5 = new ConfigChangeWatcher.ConfigChangeEvent(
              "default:\n" +
                  "  rate: 400" +
                  "  duration: abc", ConfigChangeWatcher.EventType.MODIFY);
  
          watcher.notify(value5);
          globalDefaultSamplingRateEquals(watcher, 499);
          Assertions.assertEquals(watcher.value(), "default:\n" +
              "  rate: 500");
      }
  
      public static class DefaultSampleRateMockConfigWatcherRegister extends ConfigWatcherRegister {
  
          public DefaultSampleRateMockConfigWatcherRegister(long syncPeriod) {
              super(syncPeriod);
          }
  
          @Override
          public Optional<ConfigTable> readConfig(Set<String> keys) {
              ConfigTable table = new ConfigTable();
              table.add(new ConfigTable.ConfigItem("agent-analyzer.default.traceSamplingPolicy", "default:\n" +
                  "  rate: 9000"));
              return Optional.of(table);
          }
  
          @Override
          public Optional<GroupConfigTable> readGroupConfig(final Set<String> keys) {
              return Optional.empty();
          }
      }
  
      @Test
      @Timeout(20)
      public void testServiceSampleRateDynamicUpdate() throws InterruptedException {
          ConfigWatcherRegister register = new ServiceMockConfigWatcherRegister(3);
  
          TraceSamplingPolicyWatcher watcher = new TraceSamplingPolicyWatcher(moduleConfig, provider);
          Whitebox.setInternalState(provider, "moduleConfig", moduleConfig);
          provider.getModuleConfig().setTraceSamplingPolicyWatcher(watcher);
          register.registerConfigChangeWatcher(watcher);
          register.start();
  
          while (getSamplingPolicy("serverName1", watcher) == null) {
              Thread.sleep(1000);
          }
  
          SamplingPolicy samplingPolicy = getSamplingPolicy("serverName1", watcher);
          Assertions.assertEquals(samplingPolicy.getRate().intValue(), 2000);
          Assertions.assertEquals(samplingPolicy.getDuration().intValue(), 30000);
          Assertions.assertEquals(getSamplingPolicy("serverName1", provider.getModuleConfig().getTraceSamplingPolicyWatcher())
                                  .getRate()
                                  .intValue(), 2000);
      }
  
      @Test
      public void testServiceSampleRateNotify() {
          TraceSamplingPolicyWatcher watcher = new TraceSamplingPolicyWatcher(moduleConfig, provider);
          ConfigChangeWatcher.ConfigChangeEvent value1 = new ConfigChangeWatcher.ConfigChangeEvent(
              "services:\n" +
                  "  - name: serverName1\n" +
                  "    rate: 8000\n" +
                  "    duration: 20000", ConfigChangeWatcher.EventType.MODIFY);
  
          watcher.notify(value1);
  
          Assertions.assertEquals(getSamplingPolicy("serverName1", watcher).getRate().intValue(), 8000);
          Assertions.assertEquals(getSamplingPolicy("serverName1", watcher).getDuration().intValue(), 20000);
          Assertions.assertEquals(watcher.value(), "services:\n" +
              "  - name: serverName1\n" +
              "    rate: 8000\n" +
              "    duration: 20000");
  
          // use serverName1's sampling rate
          Assertions.assertTrue(watcher.shouldSample("serverName1", 7999, -1));
          Assertions.assertTrue(watcher.shouldSample("serverName1", 10000, 20000));
  
          ConfigChangeWatcher.ConfigChangeEvent value2 = new ConfigChangeWatcher.ConfigChangeEvent(
              "", ConfigChangeWatcher.EventType.DELETE);
  
          watcher.notify(value2);
  
          Assertions.assertNull(getSamplingPolicy("serverName1", watcher));
          // use global sampling rate
          Assertions.assertTrue(watcher.shouldSample("serverName1", 9999, -1));
          Assertions.assertFalse(watcher.shouldSample("serverName1", 10000, 1));
  
          Assertions.assertEquals(watcher.value(), null);
  
          ConfigChangeWatcher.ConfigChangeEvent value3 = new ConfigChangeWatcher.ConfigChangeEvent(
              "services:\n" +
                  "  - name: serverName1\n" +
                  "    rate: 8000\n" +
                  "    duration: 20000", ConfigChangeWatcher.EventType.ADD);
  
          watcher.notify(value3);
          Assertions.assertEquals(getSamplingPolicy("serverName1", watcher).getRate().intValue(), 8000);
          Assertions.assertEquals(getSamplingPolicy("serverName1", watcher).getDuration().intValue(), 20000);
          Assertions.assertTrue(watcher.shouldSample("serverName1", 7999, -1));
          Assertions.assertTrue(watcher.shouldSample("serverName1", 10000, 20000));
  
          Assertions.assertEquals(watcher.value(), "services:\n" +
              "  - name: serverName1\n" +
              "    rate: 8000\n" +
              "    duration: 20000");
  
          ConfigChangeWatcher.ConfigChangeEvent value4 = new ConfigChangeWatcher.ConfigChangeEvent(
              "services:\n" +
                  "  - name: serverName1\n" +
                  "    rate: 9000\n" +
                  "    duration: 30000", ConfigChangeWatcher.EventType.MODIFY);
  
          watcher.notify(value4);
          Assertions.assertEquals(getSamplingPolicy("serverName1", watcher).getRate().intValue(), 9000);
          Assertions.assertEquals(getSamplingPolicy("serverName1", watcher).getDuration().intValue(), 30000);
          Assertions.assertTrue(watcher.shouldSample("serverName1", 8999, -1));
          Assertions.assertTrue(watcher.shouldSample("serverName1", 10000, 30000));
  
          Assertions.assertEquals(watcher.value(), "services:\n" +
              "  - name: serverName1\n" +
              "    rate: 9000\n" +
              "    duration: 30000");
  
          ConfigChangeWatcher.ConfigChangeEvent value5 = new ConfigChangeWatcher.ConfigChangeEvent(
              "services:\n" +
                  "  - name: serverName1\n" +
                  "    rate: 8000\n" +
                  "    duration: abc", ConfigChangeWatcher.EventType.MODIFY);
  
          watcher.notify(value5);
          Assertions.assertEquals(getSamplingPolicy("serverName1", watcher).getRate().intValue(), 9000);
          Assertions.assertEquals(getSamplingPolicy("serverName1", watcher).getDuration().intValue(), 30000);
          Assertions.assertTrue(watcher.shouldSample("serverName1", 8999, -1));
          Assertions.assertTrue(watcher.shouldSample("serverName1", 10000, 30000));
  
          Assertions.assertEquals(watcher.value(), "services:\n" +
              "  - name: serverName1\n" +
              "    rate: 9000\n" +
              "    duration: 30000");
  
      }
  
      public static class ServiceMockConfigWatcherRegister extends ConfigWatcherRegister {
  
          public ServiceMockConfigWatcherRegister(long syncPeriod) {
              super(syncPeriod);
          }
  
          @Override
          public Optional<ConfigTable> readConfig(Set<String> keys) {
              ConfigTable table = new ConfigTable();
              table.add(new ConfigTable.ConfigItem("agent-analyzer.default.traceSamplingPolicy", "services:\n" +
                  "  - name: serverName1\n" +
                  "    rate: 2000\n" +
                  "    duration: 30000"));
              return Optional.of(table);
          }
  
          @Override
          public Optional<GroupConfigTable> readGroupConfig(final Set<String> keys) {
              return Optional.empty();
          }
      }
  
      private void globalDefaultSamplingRateEquals(TraceSamplingPolicyWatcher watcher, int sample) {
          Assertions.assertTrue(watcher.shouldSample("", sample, -1));
          Assertions.assertFalse(watcher.shouldSample("", sample + 1, -1));
      }
  
      private void globalDefaultDurationEquals(TraceSamplingPolicyWatcher watcher, int duration) {
          Assertions.assertTrue(watcher.shouldSample("", 10000, duration));
          Assertions.assertFalse(watcher.shouldSample("", 10000, duration - 1));
      }
  
      private SamplingPolicy getSamplingPolicy(String service, TraceSamplingPolicyWatcher watcher) {
          AtomicReference<SamplingPolicySettings> samplingPolicySettings = Whitebox.getInternalState(
              watcher, "samplingPolicySettings");
          return samplingPolicySettings.get().get(service);
      }
  }
  