From a6bdbf0b0bcdb17d754db775912bea18e14440d6 Mon Sep 17 00:00:00 2001
From: Martijn van Groningen <martijn.v.groningen@gmail.com>
Date: Thu, 19 Mar 2015 13:08:47 -0700
Subject: [PATCH] Added simple benchmark to test fired alert throughput

Original commit: elastic/x-pack-elasticsearch@186b6305ff7c35074e06c34d59a214e9945269d4
---
 .../alerts/test/bench/AlertsBenchmark.java    | 135 ++++++++++++++++++
 1 file changed, 135 insertions(+)
 create mode 100644 src/test/java/org/elasticsearch/alerts/test/bench/AlertsBenchmark.java

diff --git a/src/test/java/org/elasticsearch/alerts/test/bench/AlertsBenchmark.java b/src/test/java/org/elasticsearch/alerts/test/bench/AlertsBenchmark.java
new file mode 100644
index 00000000000..2f3fdad7dc5
--- /dev/null
+++ b/src/test/java/org/elasticsearch/alerts/test/bench/AlertsBenchmark.java
@@ -0,0 +1,135 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+package org.elasticsearch.alerts.test.bench;
+
+import org.elasticsearch.action.search.SearchRequest;
+import org.elasticsearch.alerts.AlertsPlugin;
+import org.elasticsearch.alerts.client.AlertSourceBuilder;
+import org.elasticsearch.alerts.client.AlertsClient;
+import org.elasticsearch.alerts.scheduler.Scheduler;
+import org.elasticsearch.alerts.scheduler.SchedulerMock;
+import org.elasticsearch.alerts.scheduler.SchedulerModule;
+import org.elasticsearch.alerts.transport.actions.put.PutAlertRequest;
+import org.elasticsearch.common.collect.ImmutableList;
+import org.elasticsearch.common.inject.Module;
+import org.elasticsearch.common.logging.Loggers;
+import org.elasticsearch.common.settings.ImmutableSettings;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.xcontent.XContentType;
+import org.elasticsearch.node.NodeBuilder;
+import org.elasticsearch.node.internal.InternalNode;
+import org.elasticsearch.search.builder.SearchSourceBuilder;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+
+import static org.elasticsearch.alerts.actions.ActionBuilders.indexAction;
+import static org.elasticsearch.alerts.condition.ConditionBuilders.scriptCondition;
+import static org.elasticsearch.alerts.input.InputBuilders.searchInput;
+import static org.elasticsearch.alerts.scheduler.schedule.Schedules.interval;
+
+/**
+ */
+public class AlertsBenchmark {
+
+    public static void main(String[] args) throws Exception {
+        Settings settings = ImmutableSettings.builder()
+                .put("plugins.load_classpath_plugins", false)
+                .put("plugin.types", BenchmarkAlertPlugin.class.getName())
+                .put("cluster.name", AlertsBenchmark.class.getSimpleName())
+                .put("http.cors.enabled", true)
+                .build();
+        InternalNode node = (InternalNode) NodeBuilder.nodeBuilder().settings(settings).data(false).node();
+        node.client().admin().cluster().prepareHealth().setWaitForGreenStatus().get();
+        Thread.sleep(5000);
+        AlertsClient alertsClient = node.injector().getInstance(AlertsClient.class);
+
+        int numAlerts = 1000;
+        for (int i = 0; i < numAlerts; i++) {
+            final String name = "_name" + i;
+            PutAlertRequest putAlertRequest = new PutAlertRequest(
+                    new AlertSourceBuilder()
+                            .schedule(interval("5s"))
+                            .input(searchInput(new SearchRequest().source(
+                                    new SearchSourceBuilder()
+                            )))
+                            .condition(scriptCondition("1 == 1"))
+                            .addAction(indexAction("index", "type"))
+                            .buildAsBytes(XContentType.JSON)
+            );
+            putAlertRequest.setAlertName(name);
+            alertsClient.putAlert(putAlertRequest).actionGet();
+        }
+
+        int numThreads = 50;
+        int alertsPerThread = numAlerts / numThreads;
+        final SchedulerMock scheduler = (SchedulerMock) node.injector().getInstance(Scheduler.class);
+        Thread[] threads = new Thread[numThreads];
+        for (int i = 0; i < numThreads; i++) {
+            final int begin = i * alertsPerThread;
+            final int end = (i + 1) * alertsPerThread;
+            Runnable r = new Runnable() {
+                @Override
+                public void run() {
+                    while (true) {
+                        for (int j = begin; j < end; j++) {
+                            scheduler.fire("_name" + j);
+                        }
+                    }
+                }
+            };
+            threads[i] = new Thread(r);
+            threads[i].start();
+        }
+
+        
+        for (Thread thread : threads) {
+            thread.join();
+        }
+    }
+
+    public static final class BenchmarkAlertPlugin extends AlertsPlugin {
+
+        public BenchmarkAlertPlugin(Settings settings) {
+            super(settings);
+            Loggers.getLogger(BenchmarkAlertPlugin.class, settings).info("using benchmark alerts plugin");
+        }
+
+        @Override
+        public Collection<Class<? extends Module>> modules() {
+            return ImmutableList.<Class<? extends Module>>of(AlertsModule.class);
+        }
+
+        public static class AlertsModule extends org.elasticsearch.alerts.AlertsModule {
+
+            @Override
+            public Iterable<? extends Module> spawnModules() {
+                List<Module> modules = new ArrayList<>();
+                for (Module module : super.spawnModules()) {
+                    if (module instanceof SchedulerModule) {
+                        // replacing scheduler module so we'll
+                        // have control on when it fires a job
+                        modules.add(new MockSchedulerModule());
+
+                    } else {
+                        modules.add(module);
+                    }
+                }
+                return modules;
+            }
+
+            public static class MockSchedulerModule extends SchedulerModule {
+
+                public MockSchedulerModule() {
+                    super(SchedulerMock.class);
+                }
+
+            }
+        }
+    }
+
+}