ParallelEnricher { public EnrichmentResult apply( JSONObject message , EnrichmentStrategies strategy , SensorEnrichmentConfig config , PerformanceLogger perfLog ) throws ExecutionException, InterruptedException { if(message == null) { return null; } if(perfLog != null) { perfLog.mark("execute"); if(perfLog.isDebugEnabled() && !cacheStats.isEmpty()) { CacheStats before = cacheStats.get(strategy); CacheStats after = concurrencyContext.getCache().stats(); if(before != null && after != null) { CacheStats delta = after.minus(before); perfLog.log("cache", delta.toString()); } cacheStats.put(strategy, after); } } String sensorType = MessageUtils.getSensorType(message); message.put(getClass().getSimpleName().toLowerCase() + ".splitter.begin.ts", "" + System.currentTimeMillis()); Map<String, List<JSONObject>> tasks = splitMessage( message , strategy , config ); message.put(getClass().getSimpleName().toLowerCase() + ".splitter.end.ts", "" + System.currentTimeMillis()); message.put(getClass().getSimpleName().toLowerCase() + ".enrich.begin.ts", "" + System.currentTimeMillis()); if(perfLog != null) { perfLog.mark("enrich"); } List<CompletableFuture<JSONObject>> taskList = new ArrayList<>(); List<Map.Entry<Object, Throwable>> errors = Collections.synchronizedList(new ArrayList<>()); for(Map.Entry<String, List<JSONObject>> task : tasks.entrySet()) { EnrichmentAdapter<CacheKey> adapter = enrichmentsByType.get(task.getKey()); if(adapter == null) { throw new IllegalStateException("Unable to find an adapter for " + task.getKey() + ", possible adapters are: " + Joiner.on(",").join(enrichmentsByType.keySet())); } message.put("adapter." + adapter.getClass().getSimpleName().toLowerCase() + ".begin.ts", "" + System.currentTimeMillis()); for(JSONObject m : task.getValue()) { for(Object o : m.keySet()) { String field = (String) o; Object value = m.get(o); if(value == null) { message.put("adapter." + adapter.getClass().getSimpleName().toLowerCase() + ".end.ts", "" + System.currentTimeMillis()); continue; } CacheKey cacheKey = new CacheKey(field, value, config); String prefix = adapter.getOutputPrefix(cacheKey); Supplier<JSONObject> supplier = () -> { try { JSONObject ret = concurrencyContext.getCache().get(cacheKey, new EnrichmentCallable(cacheKey, adapter)); if(ret == null) { ret = new JSONObject(); } JSONObject adjustedKeys = EnrichmentUtils .adjustKeys(new JSONObject(), ret, cacheKey.getField(), prefix); adjustedKeys.put("adapter." + adapter.getClass().getSimpleName().toLowerCase() + ".end.ts", "" + System.currentTimeMillis()); return adjustedKeys; } catch (Throwable e) { JSONObject errorMessage = new JSONObject(); errorMessage.putAll(m); errorMessage.put(Constants.SENSOR_TYPE, sensorType ); errors.add(new AbstractMap.SimpleEntry<>(errorMessage, new IllegalStateException(strategy + " error with " + task.getKey() + " failed: " + e.getMessage(), e))); return new JSONObject(); } }; taskList.add(CompletableFuture.supplyAsync( supplier, ConcurrencyContext.getExecutor())); } } } if(taskList.isEmpty()) { message.put(getClass().getSimpleName().toLowerCase() + ".enrich.end.ts", "" + System.currentTimeMillis()); return new EnrichmentResult(message, errors); } EnrichmentResult ret = new EnrichmentResult(all(taskList, message, (left, right) -> join(left, right)).get(), errors); ret.getResult().put(getClass().getSimpleName().toLowerCase() + ".enrich.end.ts", "" + System.currentTimeMillis()); if(perfLog != null) { String key = message.get(Constants.GUID) + ""; perfLog.log("enrich", "key={}, elapsed time to enrich", key); perfLog.log("execute", "key={}, elapsed time to run execute", key); } return ret; } ParallelEnricher( Map<String, EnrichmentAdapter<CacheKey>> enrichmentsByType
                         , ConcurrencyContext concurrencyContext
                         , boolean logStats
                         ); EnrichmentResult apply( JSONObject message
                         , EnrichmentStrategies strategy
                         , SensorEnrichmentConfig config
                         , PerformanceLogger perfLog
                         ); static CompletableFuture<JSONObject> all(
            List<CompletableFuture<JSONObject>> futures
          , JSONObject identity
          , BinaryOperator<JSONObject> reduceOp
  ); Map<String, List<JSONObject>> splitMessage( JSONObject message
                                                   , EnrichmentStrategy enrichmentStrategy
                                                   , SensorEnrichmentConfig config
                                                   );  }
@Test public void testCacheHit() throws Exception { numAccesses.set(0); JSONObject message = new JSONObject() {{ put(Constants.SENSOR_TYPE, "test"); }}; for(int i = 0;i < 10;++i) { SensorEnrichmentConfig config = JSONUtils.INSTANCE.load(goodConfig, SensorEnrichmentConfig.class); config.getConfiguration().putIfAbsent("stellarContext", stellarContext); ParallelEnricher.EnrichmentResult result = enricher.apply(message, EnrichmentStrategies.ENRICHMENT, config, null); } assertTrue(2 >= numAccesses.get()); }
@Test public void testGoodConfig() throws Exception { SensorEnrichmentConfig config = JSONUtils.INSTANCE.load(goodConfig, SensorEnrichmentConfig.class); config.getConfiguration().putIfAbsent("stellarContext", stellarContext); JSONObject message = new JSONObject() {{ put(Constants.SENSOR_TYPE, "test"); }}; ParallelEnricher.EnrichmentResult result = enricher.apply(message, EnrichmentStrategies.ENRICHMENT, config, null); JSONObject ret = result.getResult(); assertEquals(11, ret.size(), "Got the wrong result count: " + ret); assertEquals(1, ret.get("map.blah")); assertEquals("test", ret.get("source.type")); assertEquals(1, ret.get("one")); assertEquals(2, ret.get("foo")); assertEquals("TEST", ret.get("ALL_CAPS")); assertEquals(0, result.getEnrichmentErrors().size()); assertTrue(result.getResult().containsKey("adapter.accessloggingstellaradapter.begin.ts")); assertTrue(result.getResult().containsKey("adapter.accessloggingstellaradapter.end.ts")); assertTrue(result.getResult().containsKey("parallelenricher.splitter.begin.ts")); assertTrue(result.getResult().containsKey("parallelenricher.splitter.end.ts")); assertTrue(result.getResult().containsKey("parallelenricher.enrich.begin.ts")); assertTrue(result.getResult().containsKey("parallelenricher.enrich.end.ts")); }
@Test public void testNullEnrichment() throws Exception { SensorEnrichmentConfig config = JSONUtils.INSTANCE.load(nullConfig, SensorEnrichmentConfig.class); config.getConfiguration().putIfAbsent("stellarContext", stellarContext); JSONObject message = new JSONObject() {{ put(Constants.SENSOR_TYPE, "test"); }}; ParallelEnricher.EnrichmentResult result = enricher.apply(message, EnrichmentStrategies.ENRICHMENT, config, null); JSONObject ret = result.getResult(); assertEquals(7, ret.size(), "Got the wrong result count: " + ret); assertTrue(result.getResult().containsKey("adapter.dummyenrichmentadapter.begin.ts")); assertTrue(result.getResult().containsKey("adapter.dummyenrichmentadapter.end.ts")); assertTrue(result.getResult().containsKey("parallelenricher.splitter.begin.ts")); assertTrue(result.getResult().containsKey("parallelenricher.splitter.end.ts")); assertTrue(result.getResult().containsKey("parallelenricher.enrich.begin.ts")); assertTrue(result.getResult().containsKey("parallelenricher.enrich.end.ts")); }
@Test public void testBadConfig() throws Exception { SensorEnrichmentConfig config = JSONUtils.INSTANCE.load(badConfig, SensorEnrichmentConfig.class); config.getConfiguration().putIfAbsent("stellarContext", stellarContext); JSONObject message = new JSONObject() {{ put(Constants.SENSOR_TYPE, "test"); }}; ParallelEnricher.EnrichmentResult result = enricher.apply(message, EnrichmentStrategies.ENRICHMENT, config, null); JSONObject ret = result.getResult(); assertEquals(11, ret.size(), ret + " is not what I expected"); assertEquals(1, ret.get("map.blah")); assertEquals("test", ret.get("source.type")); assertEquals(1, ret.get("one")); assertEquals(2, ret.get("foo")); assertEquals("TEST", ret.get("ALL_CAPS")); assertEquals(1, result.getEnrichmentErrors().size()); assertTrue(result.getResult().containsKey("adapter.accessloggingstellaradapter.begin.ts")); assertTrue(result.getResult().containsKey("adapter.accessloggingstellaradapter.end.ts")); assertTrue(result.getResult().containsKey("parallelenricher.splitter.begin.ts")); assertTrue(result.getResult().containsKey("parallelenricher.splitter.end.ts")); assertTrue(result.getResult().containsKey("parallelenricher.enrich.begin.ts")); assertTrue(result.getResult().containsKey("parallelenricher.enrich.end.ts")); }