Examples with CacheKey org.apache.metron.enrichment.bolt.CacheKey used on opensource projects

Example 6 with CacheKey

use of org.apache.metron.enrichment.bolt.CacheKey in project metron by apache.

the class ParallelEnricher method apply.

/**
 * Fully enriches a message.  Each enrichment is done in parallel via a threadpool.
 * Each enrichment is fronted with a LRU cache.
 *
 * @param message the message to enrich
 * @param strategy The enrichment strategy to use (e.g. enrichment or threat intel)
 * @param config The sensor enrichment config
 * @param perfLog The performance logger.  We log the performance for this call, the split portion and the enrichment portion.
 * @return the enrichment result
 */
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());
    // Split the message into individual tasks.
    // 
    // A task will either correspond to an enrichment adapter or,
    // in the case of Stellar, a stellar subgroup.  The tasks will be grouped by enrichment type (the key of the
    // tasks map).  Each JSONObject will correspond to a unit of work.
    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()) {
        // task is the list of enrichment tasks for the task.getKey() adapter
        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()));
        }
        for (JSONObject m : task.getValue()) {
            /* now for each unit of work (each of these only has one element in them)
         * the key is the field name and the value is value associated with that field.
         *
         * In the case of stellar enrichment, the field name is the subgroup name or empty string.
         * The value is the subset of the message needed for the enrichment.
         *
         * In the case of another enrichment (e.g. hbase), the field name is the field name being enriched.
         * The value is the corresponding value.
         */
            for (Object o : m.keySet()) {
                String field = (String) o;
                Object value = m.get(o);
                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();
                        }
                        // each enrichment has their own unique prefix to use to adjust the keys for the enriched fields.
                        return EnrichmentUtils.adjustKeys(new JSONObject(), ret, cacheKey.getField(), prefix);
                    } 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();
                    }
                };
                // add the Future to the task list
                taskList.add(CompletableFuture.supplyAsync(supplier, ConcurrencyContext.getExecutor()));
            }
        }
    }
    if (taskList.isEmpty()) {
        return new EnrichmentResult(message, errors);
    }
    EnrichmentResult ret = new EnrichmentResult(all(taskList, message, (left, right) -> join(left, right)).get(), errors);
    message.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;
}

Example 7 with CacheKey

use of org.apache.metron.enrichment.bolt.CacheKey in project metron by apache.

the class GeoAdapterTest method testEnrich.

@Test
public void testEnrich() throws Exception {
    JSONObject actualMessage = geo.enrich(new CacheKey("dummy", IP, null));
    Assert.assertNotNull(actualMessage.get("locID"));
    Assert.assertEquals(expectedMessage, actualMessage);
}

Example 8 with CacheKey

use of org.apache.metron.enrichment.bolt.CacheKey in project metron by apache.

the class HostFromJSONListAdapterTest method testEnrichNonString.

@Test
public void testEnrichNonString() throws Exception {
    HostFromJSONListAdapter hja = new HostFromJSONListAdapter(expectedKnownHostsString);
    JSONObject actualMessage = hja.enrich(new CacheKey("dummy", ip, null));
    Assert.assertNotNull(actualMessage);
    Assert.assertEquals(expectedMessage, actualMessage);
    actualMessage = hja.enrich(new CacheKey("dummy", 10L, null));
    JSONObject emptyJson = new JSONObject();
    Assert.assertEquals(emptyJson, actualMessage);
}

Example 9 with CacheKey

use of org.apache.metron.enrichment.bolt.CacheKey in project metron by apache.

the class HostFromPropertiesFileAdapterTest method testEnrich.

@Test
public void testEnrich() throws Exception {
    Map<String, JSONObject> mapKnownHosts = new HashMap<>();
    JSONArray jsonArray = (JSONArray) JSONValue.parse(expectedKnownHostsString);
    Iterator jsonArrayIterator = jsonArray.iterator();
    while (jsonArrayIterator.hasNext()) {
        JSONObject jsonObject = (JSONObject) jsonArrayIterator.next();
        String host = (String) jsonObject.remove("ip");
        mapKnownHosts.put(host, jsonObject);
    }
    HostFromPropertiesFileAdapter hfa = new HostFromPropertiesFileAdapter(mapKnownHosts);
    JSONObject actualMessage = hfa.enrich(new CacheKey("dummy", ip, null));
    Assert.assertNotNull(actualMessage);
    Assert.assertEquals(expectedMessage, actualMessage);
    actualMessage = hfa.enrich(new CacheKey("dummy", ip1, null));
    JSONObject emptyJson = new JSONObject();
    Assert.assertEquals(emptyJson, actualMessage);
}

Example 10 with CacheKey

use of org.apache.metron.enrichment.bolt.CacheKey in project metron by apache.

the class SimpleHBaseAdapterTest method testMultiColumnFamilies.

@Test
public void testMultiColumnFamilies() throws Exception {
    SimpleHBaseAdapter sha = new SimpleHBaseAdapter();
    sha.lookup = lookup;
    SensorEnrichmentConfig broSc = JSONUtils.INSTANCE.load(sourceConfigWithCFStr, SensorEnrichmentConfig.class);
    JSONObject actualMessage = sha.enrich(new CacheKey("test", "test", broSc));
    Assert.assertEquals(actualMessage, new JSONObject());
    actualMessage = sha.enrich(new CacheKey("ip_dst_addr", "10.0.2.4", broSc));
    Assert.assertNotNull(actualMessage);
    Assert.assertEquals(new JSONObject(ImmutableMap.of("cf1.key", "value")), actualMessage);
}