001/*
002 * Licensed to the Apache Software Foundation (ASF) under one
003 * or more contributor license agreements.  See the NOTICE file
004 * distributed with this work for additional information
005 * regarding copyright ownership.  The ASF licenses this file
006 * to you under the Apache License, Version 2.0 (the
007 * "License"); you may not use this file except in compliance
008 * with the License.  You may obtain a copy of the License at
009 *
010 *     http://www.apache.org/licenses/LICENSE-2.0
011 *
012 * Unless required by applicable law or agreed to in writing, software
013 * distributed under the License is distributed on an "AS IS" BASIS,
014 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
015 * See the License for the specific language governing permissions and
016 * limitations under the License.
017 */
018package org.apache.hadoop.hbase;
019
020import static org.junit.jupiter.api.Assertions.fail;
021
022import java.lang.reflect.Constructor;
023import java.lang.reflect.Method;
024import java.time.Duration;
025import java.time.Instant;
026import java.util.Map;
027import java.util.Set;
028import java.util.concurrent.ExecutionException;
029import java.util.concurrent.ExecutorService;
030import java.util.concurrent.Executors;
031import java.util.concurrent.Future;
032import java.util.concurrent.TimeUnit;
033import java.util.concurrent.TimeoutException;
034import org.apache.hadoop.hbase.testclassification.IntegrationTests;
035import org.apache.hadoop.hbase.testclassification.LargeTests;
036import org.apache.hadoop.hbase.testclassification.MediumTests;
037import org.apache.hadoop.hbase.testclassification.SmallTests;
038import org.apache.yetus.audience.InterfaceAudience;
039import org.junit.jupiter.api.extension.AfterAllCallback;
040import org.junit.jupiter.api.extension.AfterEachCallback;
041import org.junit.jupiter.api.extension.BeforeAllCallback;
042import org.junit.jupiter.api.extension.BeforeEachCallback;
043import org.junit.jupiter.api.extension.ExtensionContext;
044import org.junit.jupiter.api.extension.ExtensionContext.Store;
045import org.junit.jupiter.api.extension.InvocationInterceptor;
046import org.junit.jupiter.api.extension.ReflectiveInvocationContext;
047import org.junit.platform.commons.JUnitException;
048import org.junit.platform.commons.util.ExceptionUtils;
049import org.slf4j.Logger;
050import org.slf4j.LoggerFactory;
051
052import org.apache.hbase.thirdparty.com.google.common.collect.ImmutableMap;
053import org.apache.hbase.thirdparty.com.google.common.collect.Iterables;
054import org.apache.hbase.thirdparty.com.google.common.collect.Sets;
055import org.apache.hbase.thirdparty.com.google.common.util.concurrent.ThreadFactoryBuilder;
056
057/**
058 * Class test rule implementation for JUnit5.
059 * <p>
060 * It ensures that all JUnit5 tests should have at least one of {@link SmallTests},
061 * {@link MediumTests}, {@link LargeTests}, {@link IntegrationTests} tags, and set timeout based on
062 * the tag.
063 * <p>
064 * It also controls the timeout for the whole test class running, while the timeout annotation in
065 * JUnit5 can only enforce the timeout for each test method. When a test is timed out, a thread dump
066 * will be printed to log output.
067 * <p>
068 * It also implements resource check for each test method, using the {@link ResourceChecker} class.
069 * <p>
070 * Finally, it also forbid System.exit call in tests. <br>
071 * TODO: need to find a new way as SecurityManager was deprecated in Java 17 and permanently
072 * disabled since Java 24.
073 */
074@InterfaceAudience.Private
075public class HBaseJupiterExtension implements InvocationInterceptor, BeforeAllCallback,
076  AfterAllCallback, BeforeEachCallback, AfterEachCallback {
077
078  private static final Logger LOG = LoggerFactory.getLogger(HBaseJupiterExtension.class);
079
080  private static final SecurityManager securityManager = new TestSecurityManager();
081
082  private static final ExtensionContext.Namespace NAMESPACE =
083    ExtensionContext.Namespace.create(HBaseJupiterExtension.class);
084
085  private static final Map<String, Duration> TAG_TO_TIMEOUT =
086    ImmutableMap.of(SmallTests.TAG, Duration.ofMinutes(3), MediumTests.TAG, Duration.ofMinutes(6),
087      LargeTests.TAG, Duration.ofMinutes(13), IntegrationTests.TAG, Duration.ZERO);
088
089  private static final String EXECUTOR = "executor";
090
091  private static final String DEADLINE = "deadline";
092
093  private static final String RESOURCE_CHECK = "rc";
094
095  private Duration pickTimeout(ExtensionContext ctx) {
096    Set<String> timeoutTags = TAG_TO_TIMEOUT.keySet();
097    Set<String> timeoutTag = Sets.intersection(timeoutTags, ctx.getTags());
098    if (timeoutTag.isEmpty()) {
099      fail("Test class " + ctx.getDisplayName() + " does not have any of the following scale tags "
100        + timeoutTags);
101    }
102    if (timeoutTag.size() > 1) {
103      fail("Test class " + ctx.getDisplayName() + " has multiple scale tags " + timeoutTag);
104    }
105    return TAG_TO_TIMEOUT.get(Iterables.getOnlyElement(timeoutTag));
106  }
107
108  @Override
109  public void beforeAll(ExtensionContext ctx) throws Exception {
110    // TODO: remove this usage
111    System.setSecurityManager(securityManager);
112    Duration timeout = pickTimeout(ctx);
113    if (timeout.isZero() || timeout.isNegative()) {
114      LOG.info("No timeout for {}", ctx.getDisplayName());
115      // zero means no timeout
116      return;
117    }
118    Instant deadline = Instant.now().plus(timeout);
119    LOG.info("Timeout for {} is {}, it should be finished before {}", ctx.getDisplayName(), timeout,
120      deadline);
121    ExecutorService executor =
122      Executors.newSingleThreadExecutor(new ThreadFactoryBuilder().setDaemon(true)
123        .setNameFormat("HBase-Test-" + ctx.getDisplayName() + "-Main-Thread").build());
124    Store store = ctx.getStore(NAMESPACE);
125    store.put(EXECUTOR, executor);
126    store.put(DEADLINE, deadline);
127  }
128
129  @Override
130  public void afterAll(ExtensionContext ctx) throws Exception {
131    Store store = ctx.getStore(NAMESPACE);
132    ExecutorService executor = store.remove(EXECUTOR, ExecutorService.class);
133    if (executor != null) {
134      executor.shutdownNow();
135    }
136    store.remove(DEADLINE);
137    // reset secutiry manager
138    System.setSecurityManager(null);
139  }
140
141  private <T> T runWithTimeout(Invocation<T> invocation, ExtensionContext ctx, String name)
142    throws Throwable {
143    Store store = ctx.getStore(NAMESPACE);
144    ExecutorService executor = store.get(EXECUTOR, ExecutorService.class);
145    if (executor == null) {
146      return invocation.proceed();
147    }
148    Instant deadline = store.get(DEADLINE, Instant.class);
149    Instant now = Instant.now();
150    if (!now.isBefore(deadline)) {
151      fail("Test " + name + " timed out, deadline is " + deadline);
152      return null;
153    }
154
155    Duration remaining = Duration.between(now, deadline);
156    LOG.info("remaining timeout for {} is {}", name, remaining);
157    Future<T> future = executor.submit(() -> {
158      try {
159        return invocation.proceed();
160      } catch (Throwable t) {
161        // follow the same pattern with junit5
162        throw ExceptionUtils.throwAsUncheckedException(t);
163      }
164    });
165    try {
166      return future.get(remaining.toNanos(), TimeUnit.NANOSECONDS);
167    } catch (InterruptedException e) {
168      Thread.currentThread().interrupt();
169      fail("Test " + name + " interrupted");
170      return null;
171    } catch (ExecutionException e) {
172      throw ExceptionUtils.throwAsUncheckedException(e.getCause());
173    } catch (TimeoutException e) {
174      printThreadDump();
175      throw new JUnitException("Test " + name + " timed out, deadline is " + deadline, e);
176    }
177  }
178
179  private void printThreadDump() {
180    LOG.info("====> TEST TIMED OUT. PRINTING THREAD DUMP. <====");
181    LOG.info(TimedOutTestsListener.buildThreadDiagnosticString());
182  }
183
184  @Override
185  public void interceptBeforeAllMethod(Invocation<Void> invocation,
186    ReflectiveInvocationContext<Method> invocationContext, ExtensionContext extensionContext)
187    throws Throwable {
188    runWithTimeout(invocation, extensionContext, extensionContext.getDisplayName() + ".beforeAll");
189  }
190
191  @Override
192  public void interceptBeforeEachMethod(Invocation<Void> invocation,
193    ReflectiveInvocationContext<Method> invocationContext, ExtensionContext extensionContext)
194    throws Throwable {
195    runWithTimeout(invocation, extensionContext, extensionContext.getDisplayName() + ".beforeEach");
196  }
197
198  @Override
199  public void interceptTestMethod(Invocation<Void> invocation,
200    ReflectiveInvocationContext<Method> invocationContext, ExtensionContext extensionContext)
201    throws Throwable {
202    runWithTimeout(invocation, extensionContext, extensionContext.getDisplayName());
203  }
204
205  @Override
206  public void interceptAfterEachMethod(Invocation<Void> invocation,
207    ReflectiveInvocationContext<Method> invocationContext, ExtensionContext extensionContext)
208    throws Throwable {
209    runWithTimeout(invocation, extensionContext, extensionContext.getDisplayName() + ".afterEach");
210  }
211
212  @Override
213  public void interceptAfterAllMethod(Invocation<Void> invocation,
214    ReflectiveInvocationContext<Method> invocationContext, ExtensionContext extensionContext)
215    throws Throwable {
216    runWithTimeout(invocation, extensionContext, extensionContext.getDisplayName() + ".afterAll");
217  }
218
219  @Override
220  public <T> T interceptTestClassConstructor(Invocation<T> invocation,
221    ReflectiveInvocationContext<Constructor<T>> invocationContext,
222    ExtensionContext extensionContext) throws Throwable {
223    return runWithTimeout(invocation, extensionContext,
224      extensionContext.getDisplayName() + ".constructor");
225  }
226
227  // below are for implementing resource checker around test method
228
229  @Override
230  public void beforeEach(ExtensionContext ctx) throws Exception {
231    ResourceChecker rc = new ResourceChecker(ctx.getDisplayName());
232    JUnitResourceCheckers.addResourceAnalyzer(rc);
233    Store store = ctx.getStore(NAMESPACE);
234    store.put(RESOURCE_CHECK, rc);
235    rc.start();
236  }
237
238  @Override
239  public void afterEach(ExtensionContext ctx) throws Exception {
240    Store store = ctx.getStore(NAMESPACE);
241    ResourceChecker rc = store.remove(RESOURCE_CHECK, ResourceChecker.class);
242    if (rc != null) {
243      rc.end();
244    }
245  }
246}