Skip to content

Commit

Permalink
Reducing ClassGraph scans to a single scan (#2253)
Browse files Browse the repository at this point in the history
Co-authored-by: Aaron Klish <[email protected]>
  • Loading branch information
aklish and Aaron Klish authored Aug 13, 2021
1 parent 6526099 commit 8e15a2d
Show file tree
Hide file tree
Showing 4 changed files with 88 additions and 30 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
import lombok.Getter;

import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedHashMap;
Expand All @@ -32,16 +33,13 @@
public class HashMapDataStore implements DataStore, DataStoreTestHarness {
protected final Map<Type<?>, Map<String, Object>> dataStore = Collections.synchronizedMap(new HashMap<>());
@Getter protected EntityDictionary dictionary;
@Getter private final Set<Package> beanPackages;
@Getter private final ConcurrentHashMap<Type<?>, AtomicLong> typeIds = new ConcurrentHashMap<>();

public HashMapDataStore(Package beanPackage) {
this(Sets.newHashSet(beanPackage));
}

public HashMapDataStore(Set<Package> beanPackages) {
this.beanPackages = beanPackages;

for (Package beanPackage : beanPackages) {
ClassScanner.getAllClasses(beanPackage.getName()).stream()
.map(ClassType::new)
Expand All @@ -52,6 +50,15 @@ public HashMapDataStore(Set<Package> beanPackages) {
}
}

public HashMapDataStore(Collection<Class<?>> beanClasses) {
beanClasses.stream()
.map(ClassType::new)
.filter(modelType -> dictionary.getFirstAnnotation(modelType,
Arrays.asList(Include.class, Exclude.class)) instanceof Include)
.forEach(modelType -> dataStore.put(modelType,
Collections.synchronizedMap(new LinkedHashMap<>())));
}

@Override
public void populateEntityDictionary(EntityDictionary dictionary) {
for (Type<?> clazz : dataStore.keySet()) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,15 +11,50 @@

import java.lang.annotation.Annotation;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

/**
* Scans a package for classes by looking at files in the classpath.
*/
public class ClassScanner {

/**
* Class Scanning is terribly costly for service boot, so we do all the scanning up front once to
* save on startup costs. All Annotations Elide scans for must be listed here:
*/
private static final String [] CACHE_ANNOTATIONS = {
//Elide Core Annotations
"com.yahoo.elide.annotation.Include",
"com.yahoo.elide.annotation.SecurityCheck",
"com.yahoo.elide.core.utils.coerce.converters.ElideTypeConverter",

//Aggregation Store Annotations
"com.yahoo.elide.datastores.aggregation.queryengines.sql.annotation.FromTable",
"com.yahoo.elide.datastores.aggregation.queryengines.sql.annotation.FromSubquery",
"org.hibernate.annotations.Subselect",
"javax.persistence.Entity",
"javax.persistence.Table"
};

private static final Map<String, Set<Class<?>>> STARTUP_CACHE = new HashMap<>();

static {
try (ScanResult scanResult = new ClassGraph().enableClassInfo().enableAnnotationInfo().scan()) {
for (String annotationName : CACHE_ANNOTATIONS) {
STARTUP_CACHE.put(annotationName, scanResult.getClassesWithAnnotation(annotationName)
.stream()
.map(ClassInfo::loadClass)
.collect(Collectors.toSet()));

}
}
}

/**
* Scans all classes accessible from the context class loader which belong to the given package and subpackages.
*
Expand All @@ -39,12 +74,11 @@ static public Set<Class<?>> getAnnotatedClasses(Package toScan, Class<? extends
* @return The classes
*/
static public Set<Class<?>> getAnnotatedClasses(String packageName, Class<? extends Annotation> annotation) {
try (ScanResult scanResult = new ClassGraph()
.enableClassInfo().enableAnnotationInfo().whitelistPackages(packageName).scan()) {
return scanResult.getClassesWithAnnotation(annotation.getCanonicalName()).stream()
.map((ClassInfo::loadClass))
.collect(Collectors.toSet());
}
return STARTUP_CACHE.get(annotation.getCanonicalName()).stream()
.filter(clazz ->
clazz.getPackage().getName().equals(packageName)
|| clazz.getPackage().getName().startsWith(packageName + "."))
.collect(Collectors.toSet());
}

/**
Expand All @@ -57,14 +91,13 @@ static public Set<Class<?>> getAnnotatedClasses(String packageName, Class<? exte
static public Set<Class<?>> getAnnotatedClasses(List<Class<? extends Annotation>> annotations,
FilterExpression filter) {
Set<Class<?>> result = new HashSet<>();
try (ScanResult scanResult = new ClassGraph().enableClassInfo().enableAnnotationInfo().scan()) {
for (Class<? extends Annotation> annotation : annotations) {
result.addAll(scanResult.getClassesWithAnnotation(annotation.getCanonicalName()).stream()
.map(ClassInfo::loadClass)
.filter(filter::include)
.collect(Collectors.toSet()));
}

for (Class<? extends Annotation> annotation : annotations) {
result.addAll(STARTUP_CACHE.get(annotation.getCanonicalName()).stream()
.filter(filter::include)
.collect(Collectors.toSet()));
}

return result;
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,11 +7,11 @@

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;
import com.yahoo.elide.annotation.ReadPermission;
import com.yahoo.elide.annotation.UpdatePermission;
import com.yahoo.elide.annotation.Include;
import org.junit.jupiter.api.Test;

import java.util.Set;
import javax.persistence.Entity;

public class ClassScannerTest {

Expand All @@ -24,25 +24,25 @@ public void testGetAllClasses() {

@Test
public void testGetAnnotatedClasses() {
Set<Class<?>> classes = ClassScanner.getAnnotatedClasses("example", ReadPermission.class);
assertEquals(6, classes.size(), "Actual: " + classes);
classes.forEach(cls -> assertTrue(cls.isAnnotationPresent(ReadPermission.class)));
Set<Class<?>> classes = ClassScanner.getAnnotatedClasses("example", Include.class);
assertEquals(30, classes.size(), "Actual: " + classes);
classes.forEach(cls -> assertTrue(cls.isAnnotationPresent(Include.class)));
}

@Test
public void testGetAllAnnotatedClasses() {
Set<Class<?>> classes = ClassScanner.getAnnotatedClasses(ReadPermission.class);
assertEquals(12, classes.size(), "Actual: " + classes);
classes.forEach(cls -> assertTrue(cls.isAnnotationPresent(ReadPermission.class)));
Set<Class<?>> classes = ClassScanner.getAnnotatedClasses(Include.class);
assertEquals(41, classes.size(), "Actual: " + classes);
classes.forEach(cls -> assertTrue(cls.isAnnotationPresent(Include.class)));
}

@Test
public void testGetAnyAnnotatedClasses() {
Set<Class<?>> classes = ClassScanner.getAnnotatedClasses(ReadPermission.class, UpdatePermission.class);
assertEquals(17, classes.size());
Set<Class<?>> classes = ClassScanner.getAnnotatedClasses(Include.class, Entity.class);
assertEquals(52, classes.size());
for (Class<?> cls : classes) {
assertTrue(cls.isAnnotationPresent(ReadPermission.class)
|| cls.isAnnotationPresent(UpdatePermission.class));
assertTrue(cls.isAnnotationPresent(Include.class)
|| cls.isAnnotationPresent(Entity.class));
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -28,10 +28,14 @@
import com.yahoo.elide.datastores.aggregation.dynamic.TableType;
import com.yahoo.elide.datastores.aggregation.metadata.models.ArgumentDefinition;
import com.yahoo.elide.datastores.aggregation.metadata.models.Column;
import com.yahoo.elide.datastores.aggregation.metadata.models.Dimension;
import com.yahoo.elide.datastores.aggregation.metadata.models.Metric;
import com.yahoo.elide.datastores.aggregation.metadata.models.Namespace;
import com.yahoo.elide.datastores.aggregation.metadata.models.Table;
import com.yahoo.elide.datastores.aggregation.metadata.models.TableSource;
import com.yahoo.elide.datastores.aggregation.metadata.models.TimeDimension;
import com.yahoo.elide.datastores.aggregation.metadata.models.TimeDimensionGrain;
import com.yahoo.elide.datastores.aggregation.metadata.models.Versioned;
import com.yahoo.elide.datastores.aggregation.queryengines.sql.annotation.FromSubquery;
import com.yahoo.elide.datastores.aggregation.queryengines.sql.annotation.FromTable;
import org.apache.commons.lang3.tuple.Pair;
Expand Down Expand Up @@ -179,7 +183,21 @@ public Set<Type<?>> getDynamicTypes() {
* @param enableMetaDataStore If Enable MetaDataStore
*/
public MetaDataStore(Set<Type<?>> modelsToBind, boolean enableMetaDataStore) {
this.metadataModelClasses = ClassScanner.getAllClasses(META_DATA_PACKAGE.getName());

//Hardcoded to avoid ClassGraph scan.
this.metadataModelClasses = new HashSet<>(Arrays.asList(
Column.class,
Metric.class,
ArgumentDefinition.class,
TableSource.class,
Dimension.class,
TimeDimension.class,
TimeDimensionGrain.class,
Table.class,
Versioned.class,
Namespace.class
));

this.enableMetaDataStore = enableMetaDataStore;

modelsToBind.forEach(cls -> {
Expand Down Expand Up @@ -223,7 +241,7 @@ public void populateEntityDictionary(EntityDictionary dictionary) {

private final Function<String, HashMapDataStore> getHashMapDataStoreInitializer() {
return key -> {
HashMapDataStore hashMapDataStore = new HashMapDataStore(META_DATA_PACKAGE);
HashMapDataStore hashMapDataStore = new HashMapDataStore(metadataModelClasses);
EntityDictionary dictionary = new EntityDictionary(new HashMap<>());
metadataModelClasses.forEach(dictionary::bindEntity);
hashMapDataStore.populateEntityDictionary(dictionary);
Expand Down

0 comments on commit 8e15a2d

Please sign in to comment.