pair = resourcePairs.next();
- final String tk = checkNotNull(pair.getKey());
- final String tv = pair.getValue();
- resource.put(tk, tv);
+ while (mapEntries.hasNext()) {
+ var pair = mapEntries.next();
+ treeMap.put(checkNotNull(pair.getKey()), pair.getValue());
}
- return new Series(key, tags, resource);
+ return treeMap;
}
public static Series of(
diff --git a/heroic-component/src/main/java/com/spotify/heroic/suggest/NumSuggestionsLimit.java b/heroic-component/src/main/java/com/spotify/heroic/suggest/NumSuggestionsLimit.java
new file mode 100644
index 000000000..03176444c
--- /dev/null
+++ b/heroic-component/src/main/java/com/spotify/heroic/suggest/NumSuggestionsLimit.java
@@ -0,0 +1,101 @@
+/*
+ * Copyright (c) 2015 Spotify AB.
+ *
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package com.spotify.heroic.suggest;
+
+import com.spotify.heroic.common.OptionalLimit;
+import java.util.Optional;
+
+/**
+ * A simple class to centralize logic around limiting the number of suggestions requested from ES by
+ * Heroic. It defaults to 50 but any request to the backend (e.g. MemoryBackend) can override that
+ * number.
+ */
+public class NumSuggestionsLimit {
+
+ /** No request is allowed to request more than this many tags, keys or tag values. */
+ public static final int LIMIT_CEILING = 250;
+
+ /**
+ * How many suggestions we should request from ES, unless the suggest API request specifies
+ * otherwise.
+ *
+ * This applies to the requests made for keys, tag and tag values. This defaults to 50,
+ * otherwise * 10,000 is used as the default which is wasteful and could lag the grafana UI.
+ */
+ public static final int DEFAULT_LIMIT = 50;
+
+ private final int limit;
+
+ private NumSuggestionsLimit() {
+ limit = DEFAULT_LIMIT;
+ }
+
+ private NumSuggestionsLimit(int limit) {
+ int okLimit = Math.min(LIMIT_CEILING, limit);
+ this.limit = okLimit;
+ }
+
+ public static NumSuggestionsLimit of(Optional limit) {
+ return limit.isEmpty() ? new NumSuggestionsLimit() : new NumSuggestionsLimit(limit.get());
+ }
+
+ public static NumSuggestionsLimit of() {
+ return new NumSuggestionsLimit();
+ }
+
+ public static NumSuggestionsLimit of(int limit) {
+ return new NumSuggestionsLimit(limit);
+ }
+
+ public int getLimit() {
+ return limit;
+ }
+
+ /**
+ * use this.limit unless limit is not empty, then return the numeric result.
+ *
+ * @param limit the limit to respect if non-empty, usually from a request object
+ * @return a new NSL object - for fluent coding support
+ */
+ public NumSuggestionsLimit create(OptionalLimit limit) {
+ int num = limit.orElse(OptionalLimit.of(this.limit)).asInteger().get();
+ return new NumSuggestionsLimit(num);
+ }
+
+ /**
+ * use this.limit unless limit is not empty, then return the numeric result.
+ *
+ * @param limit the limit to respect if non-empty, usually from a request object
+ * @return the resulting, updated numeric limit
+ */
+ public int calculateNewLimit(OptionalLimit limit) {
+ return create(limit).getLimit();
+ }
+
+ public OptionalLimit asOptionalLimit() {
+ return OptionalLimit.of(limit);
+ }
+
+ public Optional asOptionalInt() {
+ return Optional.of(getLimit());
+ }
+}
diff --git a/heroic-component/src/main/java/com/spotify/heroic/suggest/SuggestBackend.java b/heroic-component/src/main/java/com/spotify/heroic/suggest/SuggestBackend.java
index b9efd8eaa..c5e7773b1 100644
--- a/heroic-component/src/main/java/com/spotify/heroic/suggest/SuggestBackend.java
+++ b/heroic-component/src/main/java/com/spotify/heroic/suggest/SuggestBackend.java
@@ -56,6 +56,7 @@ public interface SuggestBackend extends Grouped, Initializing, Collected {
AsyncFuture tagValueSuggest(TagValueSuggest.Request request);
AsyncFuture write(WriteSuggest.Request request);
+
default AsyncFuture write(WriteSuggest.Request request, Span parentSpan) {
// Ignore the parent span if the module does not specifically implement it.
return write(request);
diff --git a/heroic-component/src/test/java/com/spotify/heroic/suggest/NumSuggestionsLimitTest.java b/heroic-component/src/test/java/com/spotify/heroic/suggest/NumSuggestionsLimitTest.java
new file mode 100644
index 000000000..c151806b4
--- /dev/null
+++ b/heroic-component/src/test/java/com/spotify/heroic/suggest/NumSuggestionsLimitTest.java
@@ -0,0 +1,40 @@
+package com.spotify.heroic.suggest;
+
+import static org.junit.Assert.assertEquals;
+
+import com.spotify.heroic.common.OptionalLimit;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.runners.MockitoJUnitRunner;
+
+@RunWith(MockitoJUnitRunner.class)
+public class NumSuggestionsLimitTest {
+ public static final int LIMIT_FORTY = 40;
+ private NumSuggestionsLimit limit;
+
+ @Before
+ public void setup() {
+ this.limit = NumSuggestionsLimit.of(LIMIT_FORTY);
+ }
+
+ @Test
+ public void TestCorrectLimitIsApplied() {
+
+ // Check that a supplied limit is selected
+ int result = this.limit.calculateNewLimit(OptionalLimit.of(5));
+ assertEquals(5, result);
+
+ // Check that none of the above have affected the NSL's stored number
+ result = this.limit.calculateNewLimit(OptionalLimit.empty());
+ assertEquals(LIMIT_FORTY, result);
+
+ // Check that a giant request limit is not respected
+ result = this.limit.calculateNewLimit(OptionalLimit.of(10_000));
+ assertEquals(NumSuggestionsLimit.LIMIT_CEILING, result);
+
+ // Check that none of the above have affected the NSL's stored number. The
+ // above operations return a new object each time.
+ assertEquals(LIMIT_FORTY, this.limit.getLimit());
+ }
+}
diff --git a/heroic-dist/src/test/java/com/spotify/heroic/HeroicConfigurationTest.java b/heroic-dist/src/test/java/com/spotify/heroic/HeroicConfigurationTest.java
index ce282ab23..5dc395e8f 100644
--- a/heroic-dist/src/test/java/com/spotify/heroic/HeroicConfigurationTest.java
+++ b/heroic-dist/src/test/java/com/spotify/heroic/HeroicConfigurationTest.java
@@ -13,9 +13,11 @@
import com.spotify.heroic.lifecycle.LifeCycleNamedHook;
import com.spotify.heroic.querylogging.HttpContext;
import com.spotify.heroic.querylogging.QueryContext;
+import com.spotify.heroic.suggest.elasticsearch.SuggestBackendKV;
import com.spotify.heroic.usagetracking.UsageTracking;
import com.spotify.heroic.usagetracking.disabled.DisabledUsageTracking;
import com.spotify.heroic.usagetracking.google.GoogleAnalytics;
+import eu.toolchain.async.AsyncFuture;
import java.io.InputStream;
import java.util.List;
import java.util.Optional;
@@ -25,104 +27,70 @@
import org.junit.Test;
public class HeroicConfigurationTest {
+
+ public static final int EXPECTED_NUM_SUGGESTIONS_LIMIT = 100;
+ public static final List REFERENCE_STARTERS = ImmutableList.of(
+ "com.spotify.heroic.analytics.bigtable.BigtableMetricAnalytics",
+ "com.spotify.heroic.cluster.CoreClusterManager",
+ "com.spotify.heroic.consumer.kafka.KafkaConsumer",
+ "com.spotify.heroic.http.HttpServer",
+ "com.spotify.heroic.metadata.elasticsearch.MetadataBackendKV",
+ "com.spotify.heroic.metric.bigtable.BigtableBackend",
+ "com.spotify.heroic.metric.datastax.DatastaxBackend",
+ "com.spotify.heroic.rpc.grpc.GrpcRpcProtocolServer",
+ "com.spotify.heroic.shell.ShellServer",
+ "com.spotify.heroic.suggest.elasticsearch.SuggestBackendKV"
+ );
+ public static final List REFERENCE_STOPPERS = ImmutableList.of(
+ "com.spotify.heroic.analytics.bigtable.BigtableMetricAnalytics",
+ "com.spotify.heroic.cluster.CoreClusterManager",
+ "com.spotify.heroic.consumer.kafka.KafkaConsumer",
+ "com.spotify.heroic.http.HttpServer",
+ "com.spotify.heroic.metadata.elasticsearch.MetadataBackendKV",
+ "com.spotify.heroic.metric.bigtable.BigtableBackend",
+ "com.spotify.heroic.metric.datastax.DatastaxBackend",
+ "com.spotify.heroic.rpc.grpc.GrpcRpcProtocolServer",
+ "com.spotify.heroic.shell.ShellServer",
+ "com.spotify.heroic.suggest.elasticsearch.SuggestBackendKV"
+ );
+
@Test
public void testMetricsLimits() throws Exception {
testConfiguration("heroic-metrics-limits.yml");
}
@Test
- public void testAll() throws Exception {
- final List referenceStarters = ImmutableList.of(
- "com.spotify.heroic.analytics.bigtable.BigtableMetricAnalytics",
- "com.spotify.heroic.cluster.CoreClusterManager",
- "com.spotify.heroic.consumer.kafka.KafkaConsumer",
- "com.spotify.heroic.http.HttpServer",
- "com.spotify.heroic.metadata.elasticsearch.MetadataBackendKV",
- "com.spotify.heroic.metric.bigtable.BigtableBackend",
- "com.spotify.heroic.metric.datastax.DatastaxBackend",
- "com.spotify.heroic.rpc.grpc.GrpcRpcProtocolServer",
- "com.spotify.heroic.shell.ShellServer",
- "com.spotify.heroic.suggest.elasticsearch.SuggestBackendKV"
- );
+ public void testHeroicAllYaml() throws Exception {
- final List referenceStoppers = ImmutableList.of(
- "com.spotify.heroic.analytics.bigtable.BigtableMetricAnalytics",
- "com.spotify.heroic.cluster.CoreClusterManager",
- "com.spotify.heroic.consumer.kafka.KafkaConsumer",
- "com.spotify.heroic.http.HttpServer",
- "com.spotify.heroic.metadata.elasticsearch.MetadataBackendKV",
- "com.spotify.heroic.metric.bigtable.BigtableBackend",
- "com.spotify.heroic.metric.datastax.DatastaxBackend",
- "com.spotify.heroic.rpc.grpc.GrpcRpcProtocolServer",
- "com.spotify.heroic.shell.ShellServer",
- "com.spotify.heroic.suggest.elasticsearch.SuggestBackendKV"
- );
+ final var instance = testConfiguration("heroic-all.yml");
- final List referenceInternalStarters = ImmutableList.of(
- "startup future"
- );
-
- final List referenceInternalStoppers = ImmutableList.of(
- "loading executor",
- "loading scheduler"
- );
-
- final HeroicCoreInstance instance = testConfiguration("heroic-all.yml");
-
- final List starters = instance.inject(c -> {
- final CoreLifeCycleRegistry reg = (CoreLifeCycleRegistry) c.lifeCycleRegistry();
- return reg
- .starters()
- .stream()
- .map(LifeCycleNamedHook::id)
- .sorted()
- .collect(Collectors.toList());
- });
-
- final List stoppers = instance.inject(c -> {
- final CoreLifeCycleRegistry reg = (CoreLifeCycleRegistry) c.lifeCycleRegistry();
- return reg
- .stoppers()
- .stream()
- .map(LifeCycleNamedHook::id)
- .sorted()
- .collect(Collectors.toList());
- });
-
- assertEquals(referenceStarters, starters);
- assertEquals(referenceStoppers, stoppers);
-
- final List internalStarters = instance.inject(c -> {
- final CoreLifeCycleRegistry reg = (CoreLifeCycleRegistry) c.internalLifeCycleRegistry();
- return reg
- .starters()
- .stream()
- .map(LifeCycleNamedHook::id)
- .sorted()
- .collect(Collectors.toList());
- });
-
- final List internalStoppers = instance.inject(c -> {
- final CoreLifeCycleRegistry reg = (CoreLifeCycleRegistry) c.internalLifeCycleRegistry();
- return reg
- .stoppers()
- .stream()
- .map(LifeCycleNamedHook::id)
- .sorted()
- .collect(Collectors.toList());
- });
-
- assertEquals(internalStarters, referenceInternalStarters);
- assertEquals(internalStoppers, referenceInternalStoppers);
+ checkStartersAndStoppers(instance);
// Check default usage tracking settings
instance.inject(coreComponent -> {
UsageTracking tracking = coreComponent.usageTracking();
+
IsInstanceOf usageTrackingMatcher = new IsInstanceOf(GoogleAnalytics.class);
assertTrue(usageTrackingMatcher.matches(tracking));
+
return null;
});
+ // Check that the SuggestBackendKV's numSuggestionsLimit was picked up
+ // from the heroic-all.yaml config file
+ instance.inject(coreComponent -> {
+ // First, pluck out the backend from the suggest manager's list of
+ // backends
+ final var suggestBackendKV =
+ (SuggestBackendKV) coreComponent.suggestManager()
+ .groupSet().inspectAll().iterator().next().getMember();
+
+ // Then pluck out the limit we expect to see and verify it
+ int limit = suggestBackendKV.getNumSuggestionsLimit().getLimit();
+ assertEquals(EXPECTED_NUM_SUGGESTIONS_LIMIT, limit);
+
+ return null;
+ });
}
@Test
@@ -204,4 +172,52 @@ private HeroicCoreInstance testConfiguration(final String name) throws Exception
private Supplier stream(String name) {
return () -> getClass().getClassLoader().getResourceAsStream(name);
}
+
+
+ private void checkStartersAndStoppers(HeroicCoreInstance instance) throws Exception {
+ final List referenceInternalStarters = ImmutableList.of(
+ "startup future"
+ );
+
+ final List referenceInternalStoppers = ImmutableList.of(
+ "loading executor",
+ "loading scheduler"
+ );
+
+ final List starters = instance.inject(c -> {
+ final CoreLifeCycleRegistry reg = (CoreLifeCycleRegistry) c.lifeCycleRegistry();
+ return collectLifeCycleNamedHookFutures(reg.starters());
+ });
+
+ final List stoppers = instance.inject(c -> {
+ final CoreLifeCycleRegistry reg = (CoreLifeCycleRegistry) c.lifeCycleRegistry();
+ return collectLifeCycleNamedHookFutures(reg.stoppers());
+ });
+
+ assertEquals(REFERENCE_STARTERS, starters);
+ assertEquals(REFERENCE_STOPPERS, stoppers);
+
+ final List internalStarters = instance.inject(c -> {
+ final CoreLifeCycleRegistry reg = (CoreLifeCycleRegistry) c.internalLifeCycleRegistry();
+ return collectLifeCycleNamedHookFutures(reg.starters());
+ });
+
+ final List internalStoppers = instance.inject(c -> {
+ final CoreLifeCycleRegistry reg = (CoreLifeCycleRegistry) c.internalLifeCycleRegistry();
+ return collectLifeCycleNamedHookFutures(reg.stoppers());
+ });
+
+ assertEquals(internalStarters, referenceInternalStarters);
+ assertEquals(internalStoppers, referenceInternalStoppers);
+ return;
+ }
+
+ private static List collectLifeCycleNamedHookFutures(
+ List>> futures) {
+ return futures
+ .stream()
+ .map(LifeCycleNamedHook::id)
+ .sorted()
+ .collect(Collectors.toList());
+ }
}
diff --git a/heroic-dist/src/test/resources/heroic-all.yml b/heroic-dist/src/test/resources/heroic-all.yml
index b941b2e04..bada39ba2 100644
--- a/heroic-dist/src/test/resources/heroic-all.yml
+++ b/heroic-dist/src/test/resources/heroic-all.yml
@@ -13,6 +13,7 @@ metadata:
suggest:
backends:
- type: elasticsearch
+ numSuggestionsLimit: 100
backendType: kv
metrics:
diff --git a/heroic-test/src/main/java/com/spotify/heroic/test/AbstractMetricBackendIT.java b/heroic-test/src/main/java/com/spotify/heroic/test/AbstractMetricBackendIT.java
index 1f1c8eefb..be9a88ccb 100644
--- a/heroic-test/src/main/java/com/spotify/heroic/test/AbstractMetricBackendIT.java
+++ b/heroic-test/src/main/java/com/spotify/heroic/test/AbstractMetricBackendIT.java
@@ -356,12 +356,18 @@ public void testWriteAndFetchLongSeries() throws Exception {
Points points = new Points();
+ final int maxStepSize = 100_000_000; // 10^8
+
long timestamp = 1;
- long maxTimestamp = 1000000000000L;
- // timestamps [1, maxTimestamp] since we can't fetch 0 (range start is exclusive)
+ long maxTimestamp = (long) Math.pow(10, 12); // 10^12 i.e. 1 million million
+
+ // timestamps [1, maxTimestamp] since we can't fetch 0 (range start is
+ // exclusive).
+ // So `points` will end up containing a minimum of
+ // 10^12 / 10^8 = 10^4 = 10,000 Point objects.
while (timestamp < maxTimestamp) {
points.p(timestamp, random.nextDouble());
- timestamp += Math.abs(random.nextInt(100000000));
+ timestamp += Math.abs(random.nextInt(maxStepSize));
}
points.p(maxTimestamp, random.nextDouble());
diff --git a/heroic-test/src/main/java/com/spotify/heroic/test/AbstractSuggestBackendIT.java b/heroic-test/src/main/java/com/spotify/heroic/test/AbstractSuggestBackendIT.java
index 25c9d41f9..521efcae6 100644
--- a/heroic-test/src/main/java/com/spotify/heroic/test/AbstractSuggestBackendIT.java
+++ b/heroic-test/src/main/java/com/spotify/heroic/test/AbstractSuggestBackendIT.java
@@ -27,7 +27,6 @@
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.ImmutableSortedSet;
import com.spotify.heroic.HeroicConfig;
import com.spotify.heroic.HeroicCore;
import com.spotify.heroic.HeroicCoreInstance;
@@ -38,28 +37,36 @@
import com.spotify.heroic.dagger.LoadingComponent;
import com.spotify.heroic.filter.TrueFilter;
import com.spotify.heroic.suggest.KeySuggest;
+import com.spotify.heroic.suggest.KeySuggest.Suggestion;
import com.spotify.heroic.suggest.MatchOptions;
+import com.spotify.heroic.suggest.NumSuggestionsLimit;
import com.spotify.heroic.suggest.SuggestBackend;
import com.spotify.heroic.suggest.SuggestManagerModule;
import com.spotify.heroic.suggest.SuggestModule;
import com.spotify.heroic.suggest.TagKeyCount;
import com.spotify.heroic.suggest.TagSuggest;
+import com.spotify.heroic.suggest.TagSuggest.Request;
import com.spotify.heroic.suggest.TagValueSuggest;
import com.spotify.heroic.suggest.TagValuesSuggest;
import com.spotify.heroic.suggest.WriteSuggest;
+import com.spotify.heroic.test.TimestampPrepender.EntityType;
import eu.toolchain.async.AsyncFramework;
import eu.toolchain.async.AsyncFuture;
import eu.toolchain.async.RetryPolicy;
+import java.time.Instant;
import java.util.ArrayList;
+import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.Set;
+import java.util.TreeSet;
import java.util.UUID;
import java.util.concurrent.ExecutionException;
import java.util.stream.Collectors;
import org.apache.commons.lang3.tuple.ImmutablePair;
import org.apache.commons.lang3.tuple.Pair;
+import org.jetbrains.annotations.NotNull;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;
@@ -68,79 +75,53 @@
@RunWith(MockitoJUnitRunner.class)
public abstract class AbstractSuggestBackendIT {
- protected final String testName = "heroic-it-" + UUID.randomUUID().toString();
-
- private final Series s1 = Series.of("aa1", ImmutableMap.of("role", "foo"));
- private final Series s2 = Series.of("aa2", ImmutableMap.of("role", "bar"));
- private final Series s3 = Series.of("bb3", ImmutableMap.of("role", "baz"));
-
- protected final DateRange range = new DateRange(0L, 0L);
-
- protected final List> testSeries =
- new ArrayList<>() {
- {
- add(new ImmutablePair<>(s1, range));
- add(new ImmutablePair<>(s2, range));
- add(new ImmutablePair<>(s3, range));
- }
- };
- private final TagValuesSuggest.Request tagValuesSuggestReq =
- new TagValuesSuggest.Request(TrueFilter.get(), range, OptionalLimit.empty(),
- OptionalLimit.empty(), ImmutableList.of());
-
- private final TagValueSuggest.Request tagValueSuggestReq =
- new TagValueSuggest.Request(TrueFilter.get(), range, OptionalLimit.empty(),
- Optional.of("role"));
-
- private final TagKeyCount.Request tagKeyCountReq =
- new TagKeyCount.Request(TrueFilter.get(), range, OptionalLimit.empty(),
- OptionalLimit.empty());
-
- private final TagSuggest.Request tagSuggestReq =
- new TagSuggest.Request(TrueFilter.get(), range, OptionalLimit.empty(),
- MatchOptions.builder().build(), Optional.empty(), Optional.of("ba"));
-
- private final KeySuggest.Request keySuggestReq =
- new KeySuggest.Request(TrueFilter.get(), range, OptionalLimit.empty(),
- MatchOptions.builder().build(), Optional.of("aa"));
+ // The requests will either not specify a limit or specify one of fifteen.
+ public static final int REQ_SUGGESTION_ENTITY_LIMIT = 15;
+ public static final String STARTS_WITH_RO = "ro"; // e.g. role
+ public static final int EFFECTIVELY_NO_LIMIT = 100_000;
+
+ private static final int SMALL_SERIES_SIZE = 3;
+ private static final int LARGE_NUM_ENTITIES = 20;
+ private static final int VERY_LARGE_NUM_ENTITIES = 500;
+ public static final String BAR = "bar";
+ public static final String BAZ = "baz";
+ public static final String FOO = "foo";
+ public static final String AA_2 = "aa2";
+ public static final String AA = "aa";
+ public static final String ROLE = "role";
+ public static final String AA_1 = "aa1";
+ public static final String BB_3 = "bb3";
+ protected final String testName = "heroic-it-" + UUID.randomUUID().toString();
- private HeroicCoreInstance core;
+ // MetaData and Suggest have no concept of datetime ranges so just set
+ // the same for all.
+ protected static final DateRange UNIVERSAL_RANGE = new DateRange(0L, 0L);
protected AsyncFramework async;
protected SuggestBackend backend;
+ private HeroicCoreInstance core;
+
protected abstract SuggestModule setupModule() throws Exception;
@Before
public final void abstractSetup() throws Exception {
- final HeroicConfig.Builder fragment = HeroicConfig
- .builder()
- .suggest(SuggestManagerModule.builder()
- .backends(ImmutableList.of(setupModule())));
-
- core = HeroicCore
- .builder()
- .setupService(false)
- .setupShellServer(false)
- .configFragment(fragment)
- .build()
+ final HeroicConfig.Builder fragment = HeroicConfig.builder()
+ .suggest(SuggestManagerModule.builder().backends(ImmutableList.of(setupModule())));
+
+ core = HeroicCore.builder().setupService(false).setupShellServer(false)
+ .configFragment(fragment).build()
.newInstance();
core.start().get();
async = core.inject(LoadingComponent::async);
- backend = core
- .inject(c -> c
- .suggestManager()
- .groupSet()
- .inspectAll()
- .stream()
- .map(GroupMember::getMember)
- .findFirst())
- .orElseThrow(() -> new IllegalStateException("Failed to find backend"));
-
+ backend = core.inject(
+ c -> c.suggestManager().groupSet().inspectAll().stream()
+ .map(GroupMember::getMember).findFirst()).orElseThrow(
+ () -> new IllegalStateException("Failed to find backend"));
}
@After
@@ -148,74 +129,205 @@ public final void abstractTeardown() throws Exception {
core.shutdown().get();
}
+ @Test
+ public void tagValuesSuggestSmall() throws Exception {
+ // Check a single suggestion with values
+ final long timestamp = getUniqueTimestamp();
+ writeSeries(backend, createSmallSeries(timestamp, EntityType.TAG));
- @Test
- public void tagValuesSuggest() throws Exception {
- writeSeries(backend, testSeries);
+ var result = getTagValuesSuggest(
+ buildTagValuesRequest(OptionalLimit.empty()));
+ var suggestion = result.getSuggestions().get(0);
- final TagValuesSuggest result = getTagValuesSuggest(tagValuesSuggestReq);
- final TagValuesSuggest.Suggestion s = result.getSuggestions().get(0);
+ var expected = new TreeSet(Arrays.asList(BAR, BAZ, FOO));
- assertEquals(
- new TagValuesSuggest.Suggestion("role", ImmutableSortedSet.of("bar", "baz", "foo"),
- false), s);
+ assertEquals(new TagValuesSuggest.Suggestion(TimestampPrepender.prepend(timestamp,
+ ROLE), expected, false), suggestion);
}
@Test
- public void tagValueSuggest() throws Exception {
- writeSeries(backend, testSeries);
+ public void tagValuesTruncatedSuggest() throws Exception {
+
+ // Check that a number of tag values larger than the supplied limit is
+ // correctly truncated.
+ final long timestamp = getUniqueTimestamp();
- final TagValueSuggest result = getTagValueSuggest(tagValueSuggestReq);
+ var largeNumTagsSeries =
+ createTestSeriesData(1, LARGE_NUM_ENTITIES, timestamp, EntityType.TAG);
+ writeSeries(backend, largeNumTagsSeries);
- assertEquals(ImmutableSet.of("bar", "baz", "foo"), ImmutableSet.copyOf(result.getValues()));
+ var result =
+ getTagValuesSuggest(
+ buildTagValuesRequest(OptionalLimit.of(REQ_SUGGESTION_ENTITY_LIMIT)));
+
+ final var suggestions = result.getSuggestions();
+ assertEquals(1, suggestions.size());
+ assertEquals(REQ_SUGGESTION_ENTITY_LIMIT, suggestions.get(0).getValues().size());
}
@Test
public void tagKeyCount() throws Exception {
- writeSeries(backend, testSeries);
+ final long timestamp = getUniqueTimestamp();
- final TagKeyCount result = getTagKeyCount(tagKeyCountReq);
+ var smallTestSeries = createSmallSeries(timestamp, EntityType.TAG);
+ writeSeries(backend, smallTestSeries);
+
+ final TagKeyCount result = getTagKeyCount(createTagCountRequest(timestamp));
final TagKeyCount.Suggestion s = result.getSuggestions().get(0);
- assertEquals("role", s.getKey());
+ assertEquals(TimestampPrepender.prepend(timestamp, ROLE), s.getKey());
assertEquals(3, s.getCount());
}
+ /**
+ * Check we get the expected tag and 3 results
+ */
@Test
- public void tagSuggest() throws Exception {
- writeSeries(backend, testSeries);
+ public void tagSuggestSmall() throws Exception {
- final Set> result = getTagSuggest(tagSuggestReq);
+ final long timestamp = getUniqueTimestamp();
+ var smallTestSeries =
+ createSmallSeries(timestamp, EntityType.TAG);
+ writeSeries(backend, smallTestSeries); // adds 3 tags
- assertEquals(ImmutableSet.of(Pair.of("role", "bar"), Pair.of("role", "baz")), result);
+ var result = getTagSuggest(
+ buildTagSuggestRequest(STARTS_WITH_RO, timestamp));
+
+ assertEquals(SMALL_SERIES_SIZE, result.size());
+ assertEquals(TimestampPrepender.prepend(timestamp, ROLE),
+ result.stream().findFirst().get().getKey());
}
+
+ /**
+ * Check that a request limit is respected and one without gets the whole lot.
+ */
+ @Test
+ public void tagSuggestLimit() throws Exception {
+
+ long timestamp = getUniqueTimestamp();
+
+ // add LARGE_NUM_ENTITIES tags. Total is now 23
+ var largeNumTagsSeries =
+ createTestSeriesData(1, LARGE_NUM_ENTITIES, timestamp, EntityType.TAG);
+ writeSeries(backend, largeNumTagsSeries);
+
+ var result =
+ getTagSuggest(buildTagSuggestRequest(STARTS_WITH_RO, timestamp,
+ REQ_SUGGESTION_ENTITY_LIMIT));
+
+ assertEquals(REQ_SUGGESTION_ENTITY_LIMIT, result.size());
+
+ // Check that the request without a limit returns the whole lot. Note that
+ // the maximum number of tags for a key is LARGE_NUM_ENTITIES - see
+ // createTestSeriesData.
+ result = getTagSuggest(buildTagSuggestRequest(STARTS_WITH_RO, timestamp));
+ assertEquals(LARGE_NUM_ENTITIES, result.size());
+ }
+
+ /**
+ * Check that a hard ceiling of NumSuggestionsLimit.LIMIT_CEILING is respected
+ *
+ * @throws Exception
+ */
+ @Test
+ public void tagSuggestCeiling() throws Exception {
+
+ long timestamp = getUniqueTimestamp();
+ var veryLargeNumTagsSeries =
+ createTestSeriesData(1, VERY_LARGE_NUM_ENTITIES, timestamp, EntityType.TAG);
+ writeSeries(backend, veryLargeNumTagsSeries);
+
+ var reqStartsWithRo = buildTagSuggestRequest(STARTS_WITH_RO, timestamp,
+ AbstractSuggestBackendIT.EFFECTIVELY_NO_LIMIT);
+ var result = getTagSuggest(reqStartsWithRo);
+ assertEquals(NumSuggestionsLimit.LIMIT_CEILING, result.size());
+ }
+
+ @Test
+ public void tagValueSuggestSmall() throws Exception {
+ final long timestamp = getUniqueTimestamp();
+
+ writeSeries(backend, createSmallSeries(timestamp, EntityType.TAG));
+
+ var result = getTagValueSuggest(
+ buildTagValueSuggestReq(ROLE, timestamp, OptionalLimit.empty()));
+
+ var expected = new TreeSet(Arrays.asList(BAR, BAZ, FOO));
+ assertEquals(ImmutableSet.copyOf(expected), ImmutableSet.copyOf(result.getValues()));
+ }
+
+ @Test
+ public void tagValueSuggestLimited() throws Exception {
+ final long timestamp = getUniqueTimestamp();
+
+ var largeNumTagsSeries =
+ createTestSeriesData(1, LARGE_NUM_ENTITIES, timestamp, EntityType.TAG);
+
+ writeSeries(backend, largeNumTagsSeries);
+
+ var result = getTagValueSuggest(
+ buildTagValueSuggestReq(ROLE, timestamp,
+ OptionalLimit.of(REQ_SUGGESTION_ENTITY_LIMIT)));
+
+ assertEquals(REQ_SUGGESTION_ENTITY_LIMIT, result.getValues().size());
+ }
+
+
@Test
public void keySuggest() throws Exception {
- writeSeries(backend, testSeries);
+ var et = EntityType.KEY;
+ {
+ final long timestamp = getUniqueTimestamp();
+ var smallTestSeries = createSmallSeries(timestamp, et);
+
+ writeSeries(backend, smallTestSeries);
+
+ var result = getKeySuggest(keySuggestStartsWithReq(AA, timestamp));
+ assertEquals(ImmutableSet.of(TimestampPrepender.prepend(timestamp, AA_1),
+ TimestampPrepender.prepend(timestamp,
+ AA_2)),
+ result);
+ }
+
+ {
+ final long timestamp = getUniqueTimestamp();
+
+ var largeNumKeysSeries =
+ createTestSeriesData(LARGE_NUM_ENTITIES, 1, timestamp, EntityType.KEY);
- final Set result = getKeySuggest(keySuggestReq);
+ writeSeries(backend, largeNumKeysSeries);
- assertEquals(ImmutableSet.of(s1.getKey(), s2.getKey()), result);
+ var result =
+ getKeySuggest(
+ keySuggestStartsWithReq(
+ AA, timestamp, OptionalLimit.of(REQ_SUGGESTION_ENTITY_LIMIT)));
+ assertEquals(REQ_SUGGESTION_ENTITY_LIMIT, result.size());
+ }
}
@Test
public void tagValueSuggestNoIdx() throws Exception {
- final TagValueSuggest result = getTagValueSuggest(tagValueSuggestReq);
+ final TagValueSuggest result = getTagValueSuggest(
+ buildTagValueSuggestReq(ROLE, 0L, OptionalLimit.empty()));
assertEquals(Collections.emptyList(), result.getValues());
}
@Test
public void tagValuesSuggestNoIdx() throws Exception {
- final TagValuesSuggest result = getTagValuesSuggest(tagValuesSuggestReq);
+ final TagValuesSuggest result = getTagValuesSuggest(
+ buildTagValuesRequest(OptionalLimit.empty()));
assertEquals(Collections.emptyList(), result.getSuggestions());
}
@Test
public void tagKeyCountNoIdx() throws Exception {
+ final long timestamp = getUniqueTimestamp();
+
+ var tagKeyCountReq = createTagCountRequest(timestamp);
final TagKeyCount result = getTagKeyCount(tagKeyCountReq);
assertEquals(Collections.emptyList(), result.getSuggestions());
@@ -223,58 +335,60 @@ public void tagKeyCountNoIdx() throws Exception {
@Test
public void tagSuggestNoIdx() throws Exception {
- final Set> result = getTagSuggest(tagSuggestReq);
+ final Set> result =
+ getTagSuggest(buildTagSuggestRequest("ba", getUniqueTimestamp()));
assertEquals(Collections.emptySet(), result);
}
@Test
public void keySuggestNoIdx() throws Exception {
- final Set result = getKeySuggest(keySuggestReq);
+ final Set result =
+ getKeySuggest(keySuggestStartsWithReq(AA, getUniqueTimestamp()));
assertEquals(Collections.emptySet(), result);
}
private AsyncFuture writeSeries(
- final SuggestBackend suggest, final Series s, final DateRange range
- ) throws Exception {
- return suggest
- .write(new WriteSuggest.Request(s, range))
- .lazyTransform(r -> async.retryUntilResolved(() -> checks(s),
- RetryPolicy.timed(10000, RetryPolicy.exponential(100, 200))))
+ final SuggestBackend suggest, final Series s, final DateRange range) {
+ return suggest.write(new WriteSuggest.Request(s, range)).lazyTransform(r -> async
+ .retryUntilResolved(() -> checks(s, range), RetryPolicy.timed(
+ 10000, RetryPolicy.exponential(100, 200))))
.directTransform(retry -> null);
}
- private AsyncFuture checks(final Series s) {
+ private AsyncFuture checks(final Series s, DateRange range) {
final List> checks = new ArrayList<>();
- checks.add(backend
- .tagSuggest(new TagSuggest.Request(matchKey(s.getKey()), range, OptionalLimit.empty(),
- MatchOptions.builder().build(), Optional.empty(), Optional.empty()))
+ checks.add(backend.tagSuggest(new TagSuggest.Request(
+ matchKey(s.getKey()), range,
+ OptionalLimit.empty(), MatchOptions.builder().build(),
+ Optional.empty(), Optional.empty()))
.directTransform(result -> {
if (result.getSuggestions().isEmpty()) {
- throw new IllegalStateException("No suggestion available for the given series");
+ throw new IllegalStateException("No tag suggestion available for the given "
+ + "series");
}
return null;
}));
- checks.add(backend
- .keySuggest(new KeySuggest.Request(matchKey(s.getKey()), range, OptionalLimit.empty(),
- MatchOptions.builder().build(), Optional.empty()))
- .directTransform(result -> {
- if (result.getSuggestions().isEmpty()) {
- throw new IllegalStateException("No suggestion available for the given series");
- }
+ checks.add(backend.keySuggest(new KeySuggest.Request(
+ matchKey(s.getKey()), range,
+ OptionalLimit.empty(), MatchOptions.builder().build(),
+ Optional.empty())).directTransform(result -> {
+ if (result.getSuggestions().isEmpty()) {
+ throw new IllegalStateException("No key suggestion available for the given series");
+ }
- return null;
- }));
+ return null;
+ }));
return async.collectAndDiscard(checks);
}
- private void writeSeries(final SuggestBackend backend, final List> data)
- throws Exception {
+ private void writeSeries(final SuggestBackend backend,
+ final List> data) throws Exception {
final List> writes = new ArrayList<>();
for (Pair p : data) {
@@ -300,23 +414,116 @@ private TagKeyCount getTagKeyCount(final TagKeyCount.Request req)
private Set> getTagSuggest(final TagSuggest.Request req)
throws ExecutionException, InterruptedException {
- return backend
- .tagSuggest(req)
- .get()
- .getSuggestions()
- .stream()
- .map(s -> Pair.of(s.getKey(), s.getValue()))
- .collect(Collectors.toSet());
+ return backend.tagSuggest(req).get().getSuggestions().stream()
+ .map(s -> Pair.of(s.getKey(), s.getValue())).collect(Collectors.toSet());
}
private Set getKeySuggest(final KeySuggest.Request req)
throws ExecutionException, InterruptedException {
- return backend
- .keySuggest(req)
- .get()
- .getSuggestions()
- .stream()
- .map(s -> s.getKey())
- .collect(Collectors.toSet());
+ return backend.keySuggest(req).get().getSuggestions().stream()
+ .map(Suggestion::getKey).collect(Collectors.toSet());
+ }
+
+ protected static List> createSmallSeries(long timestamp,
+ EntityType et) {
+
+ var p = new TimestampPrepender(et, timestamp);
+
+ return new ArrayList<>() {
+ {
+ add(createSeriesPair(AA_1, FOO, p));
+ add(createSeriesPair(AA_2, BAR, p));
+ add(createSeriesPair(BB_3, BAZ, p));
+ }
+
+ @NotNull
+ private ImmutablePair createSeriesPair(String key, String foo,
+ TimestampPrepender p) {
+ return new ImmutablePair<>(Series.of(
+ p.prepend(key, EntityType.KEY),
+ ImmutableMap.of(p.prepend(ROLE, EntityType.TAG),
+ p.prepend(foo, EntityType.TAG_VALUE))), UNIVERSAL_RANGE);
+ }
+ };
+ }
+
+ private static TagKeyCount.Request createTagCountRequest(long timestamp) {
+ return new TagKeyCount.Request(TrueFilter.get(),
+ new DateRange(timestamp, timestamp), OptionalLimit.empty(), OptionalLimit.empty());
+ }
+
+ @NotNull
+ private static TagValuesSuggest.Request buildTagValuesRequest(
+ OptionalLimit numSuggestionsLimit) {
+ return new TagValuesSuggest.Request(TrueFilter.get(),
+ UNIVERSAL_RANGE, numSuggestionsLimit,
+ OptionalLimit.of(EFFECTIVELY_NO_LIMIT), ImmutableList.of());
+ }
+
+ private static TagValueSuggest.Request buildTagValueSuggestReq(
+ String tagValue, long timestamp, OptionalLimit numSuggestionsLimit) {
+
+ return new TagValueSuggest.Request(TrueFilter.get(),
+ UNIVERSAL_RANGE, numSuggestionsLimit,
+ Optional.of(TimestampPrepender.prepend(timestamp, tagValue)));
+ }
+
+ @NotNull
+ private static Request buildTagSuggestRequest(String tagValue, long timestamp) {
+ return new Request(
+ TrueFilter.get(), UNIVERSAL_RANGE, OptionalLimit.empty(),
+ MatchOptions.builder().build(),
+ Optional.of(TimestampPrepender.prepend(timestamp, tagValue)),
+ Optional.empty());
+ }
+
+ @NotNull
+ private static Request buildTagSuggestRequest(
+ String tagValue, long timestamp, int numSuggestionsLimit) {
+ return new Request(TrueFilter.get(), UNIVERSAL_RANGE,
+ OptionalLimit.of(numSuggestionsLimit),
+ MatchOptions.builder().build(),
+ Optional.of(TimestampPrepender.prepend(timestamp, tagValue)), Optional.empty());
+ }
+
+ private static KeySuggest.Request keySuggestStartsWithReq(String startsWith, long timestamp) {
+ return keySuggestStartsWithReq(startsWith, timestamp, OptionalLimit.empty());
+ }
+
+ private static KeySuggest.Request keySuggestStartsWithReq(String startsWith, long timestamp,
+ OptionalLimit numSuggestionsLimit) {
+
+ return new KeySuggest.Request(TrueFilter.get(), UNIVERSAL_RANGE, numSuggestionsLimit,
+ MatchOptions.builder().build(),
+ Optional.of(TimestampPrepender.prepend(timestamp, startsWith)));
+ }
+
+ private static long getUniqueTimestamp() {
+ final long t = Instant.now().toEpochMilli() + (long) Math.random();
+ return t;
+ }
+
+
+ private static List> createTestSeriesData(int numKeys,
+ int tagsAndTagValuesPerKey, long timestamp, EntityType et) {
+
+ var p = new TimestampPrepender(et, timestamp);
+
+ var series = new ArrayList>(numKeys);
+
+ for (int i = 0; i < numKeys; i++) {
+ final var key = p.prepend(String.format(AA + "-%d", i + 1), EntityType.KEY);
+ for (int j = 0; j < tagsAndTagValuesPerKey; j++) {
+
+ final var tags =
+ ImmutableMap.of(
+ p.prepend(ROLE, EntityType.TAG),
+ p.prepend(FOO + "-" + (j + 1), EntityType.TAG_VALUE));
+
+ series.add(new ImmutablePair<>(Series.of(key, tags), UNIVERSAL_RANGE));
+ }
+ }
+
+ return series;
}
}
diff --git a/heroic-test/src/main/java/com/spotify/heroic/test/TimestampPrepender.java b/heroic-test/src/main/java/com/spotify/heroic/test/TimestampPrepender.java
new file mode 100644
index 000000000..d2a1fbaa1
--- /dev/null
+++ b/heroic-test/src/main/java/com/spotify/heroic/test/TimestampPrepender.java
@@ -0,0 +1,53 @@
+/*
+ * Copyright (c) 2015 Spotify AB.
+ *
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package com.spotify.heroic.test;
+
+/**
+ * Simple class to encapsulate conditional logic on whether a key, tag or tag value should be
+ * prepended with a [uniquely-identifying] timestamp.
+ */
+public class TimestampPrepender {
+
+ public enum EntityType {
+ KEY,
+ TAG,
+ TAG_VALUE
+ }
+
+ public TimestampPrepender(EntityType et, long timestamp) {
+ this.et = et;
+ this.timestamp = timestamp;
+ }
+
+ private EntityType et;
+ private long timestamp;
+
+ public static String prepend(long timestamp, String input) {
+ return Long.toString(timestamp) + "-" + input;
+ }
+
+ public String prepend(String input, EntityType et) {
+ return et == this.et
+ ? prepend(timestamp, input)
+ : input;
+ }
+}
diff --git a/heroic-test/src/test/java/com/spotify/heroic/suggest/elasticsearch/AbstractSuggestBackendKVIT.java b/heroic-test/src/test/java/com/spotify/heroic/suggest/elasticsearch/AbstractSuggestBackendKVIT.java
index 5f175ef55..c56c8056e 100644
--- a/heroic-test/src/test/java/com/spotify/heroic/suggest/elasticsearch/AbstractSuggestBackendKVIT.java
+++ b/heroic-test/src/test/java/com/spotify/heroic/suggest/elasticsearch/AbstractSuggestBackendKVIT.java
@@ -23,6 +23,7 @@
import static org.junit.Assert.assertEquals;
+import com.spotify.heroic.common.DateRange;
import com.spotify.heroic.elasticsearch.ClientWrapper;
import com.spotify.heroic.elasticsearch.ConnectionModule;
import com.spotify.heroic.elasticsearch.index.RotatingIndexMapping;
@@ -31,6 +32,7 @@
import com.spotify.heroic.suggest.WriteSuggest;
import com.spotify.heroic.test.AbstractSuggestBackendIT;
import com.spotify.heroic.test.ElasticSearchTestContainer;
+import com.spotify.heroic.test.TimestampPrepender.EntityType;
import org.junit.Test;
public abstract class AbstractSuggestBackendKVIT extends AbstractSuggestBackendIT {
@@ -63,10 +65,15 @@ protected SuggestModule setupModule() {
@Test
public void writeDuplicatesReturnErrorInResponse() throws Exception {
+ var smallTestSeries =
+ createSmallSeries(0L, EntityType.KEY);
+
final WriteSuggest firstWrite =
- backend.write(new WriteSuggest.Request(testSeries.get(0).getKey(), range)).get();
+ backend.write(new WriteSuggest.Request(smallTestSeries.get(0).getKey(),
+ UNIVERSAL_RANGE)).get();
final WriteSuggest secondWrite =
- backend.write(new WriteSuggest.Request(testSeries.get(0).getKey(), range)).get();
+ backend.write(new WriteSuggest.Request(smallTestSeries.get(0).getKey(),
+ UNIVERSAL_RANGE)).get();
assertEquals(0, firstWrite.getErrors().size());
assertEquals(2, secondWrite.getErrors().size());
diff --git a/heroic-test/src/test/java/com/spotify/heroic/suggest/elasticsearch/SuggestBackendKVTransportIT.java b/heroic-test/src/test/java/com/spotify/heroic/suggest/elasticsearch/SuggestBackendKVTransportIT.java
index e9b3f2d78..885b46404 100644
--- a/heroic-test/src/test/java/com/spotify/heroic/suggest/elasticsearch/SuggestBackendKVTransportIT.java
+++ b/heroic-test/src/test/java/com/spotify/heroic/suggest/elasticsearch/SuggestBackendKVTransportIT.java
@@ -5,15 +5,25 @@
import java.util.List;
public class SuggestBackendKVTransportIT extends AbstractSuggestBackendKVIT {
+
@Override
protected ClientWrapper setupClient() {
- List seeds = List.of(
- esContainer.getTcpHost().getHostName()
- + ":" + esContainer.getTcpHost().getPort());
+ List seeds =
+ List.of(
+ esContainer.getTcpHost().getHostName() + ":" + esContainer.getTcpHost().getPort());
+
+ return TransportClientWrapper.builder().clusterName("docker-cluster").seeds(seeds).build();
+ }
+
+ @Override
+ public void tagSuggestCeiling() throws Exception {
+ /* no-op because it causes:
+
+ io.grpc.Context validateGeneration
+ SEVERE: Context ancestry chain length is abnormally long. This suggests an error in application code. Length exceeded: 1000
- return TransportClientWrapper.builder()
- .clusterName("docker-cluster")
- .seeds(seeds)
- .build();
+ Also, my understanding is that *Transport* is to be deprecated,
+ hence there's no point spending effort trying to get this working.
+ */
}
}
diff --git a/suggest/elasticsearch/src/main/java/com/spotify/heroic/suggest/elasticsearch/ElasticsearchSuggestModule.java b/suggest/elasticsearch/src/main/java/com/spotify/heroic/suggest/elasticsearch/ElasticsearchSuggestModule.java
index 38ce8ca1b..022c68608 100644
--- a/suggest/elasticsearch/src/main/java/com/spotify/heroic/suggest/elasticsearch/ElasticsearchSuggestModule.java
+++ b/suggest/elasticsearch/src/main/java/com/spotify/heroic/suggest/elasticsearch/ElasticsearchSuggestModule.java
@@ -53,6 +53,7 @@
import com.spotify.heroic.lifecycle.LifeCycle;
import com.spotify.heroic.lifecycle.LifeCycleManager;
import com.spotify.heroic.statistics.HeroicReporter;
+import com.spotify.heroic.suggest.NumSuggestionsLimit;
import com.spotify.heroic.suggest.SuggestBackend;
import com.spotify.heroic.suggest.SuggestModule;
import dagger.Component;
@@ -94,6 +95,7 @@ public final class ElasticsearchSuggestModule implements SuggestModule, DynamicM
private final String distributedCacheSrvRecord;
private final String templateName;
private final boolean configure;
+ private final NumSuggestionsLimit numSuggestionsLimit;
private static Supplier defaultSetup = SuggestBackendKV.factory();
@@ -123,7 +125,8 @@ public ElasticsearchSuggestModule(
@JsonProperty("distributedCacheSrvRecord") Optional distributedCacheSrvRecord,
@JsonProperty("templateName") Optional templateName,
@JsonProperty("backendType") Optional backendType,
- @JsonProperty("configure") Optional configure
+ @JsonProperty("configure") Optional configure,
+ @JsonProperty("numSuggestionsLimit") Optional numSuggestionsLimit
) {
this.id = id;
this.groups = groups.orElseGet(Groups::empty).or(DEFAULT_GROUP);
@@ -142,6 +145,8 @@ public ElasticsearchSuggestModule(
this.templateName = templateName.orElse(DEFAULT_TEMPLATE_NAME);
this.type = backendType.map(this::lookupBackendType).orElse(defaultSetup);
this.configure = configure.orElse(DEFAULT_CONFIGURE);
+
+ this.numSuggestionsLimit = NumSuggestionsLimit.of(numSuggestionsLimit);
}
private Supplier lookupBackendType(final String bt) {
@@ -156,21 +161,24 @@ private Supplier lookupBackendType(final String bt) {
}
@Override
- public Exposed module(PrimaryComponent primary, Depends depends, final String id) {
+ public Exposed module(PrimaryComponent primary,
+ Depends depends,
+ final String id) {
final BackendType backendType = type.get();
return DaggerElasticsearchSuggestModule_C
- .builder()
- .primaryComponent(primary)
- .depends(depends)
- .connectionModule(connection)
- .m(new M(backendType))
- .build();
+ .builder()
+ .primaryComponent(primary)
+ .depends(depends)
+ .connectionModule(connection)
+ .m(new M(backendType))
+ .o(new O())
+ .build();
}
@ElasticsearchScope
- @Component(modules = {M.class, ConnectionModule.class},
- dependencies = {PrimaryComponent.class, Depends.class})
+ @Component(modules = {M.class, O.class, ConnectionModule.class},
+ dependencies = {PrimaryComponent.class, Depends.class})
interface C extends Exposed {
@Override
SuggestBackend backend();
@@ -207,6 +215,13 @@ public boolean configure(ExtraParameters params) {
params.contains(ELASTICSEARCH_CONFIGURE_PARAM);
}
+ @Provides
+ @ElasticsearchScope
+ @Named("numSuggestionsLimit")
+ public int numSuggestionsLimit(ExtraParameters params) {
+ return numSuggestionsLimit.getLimit();
+ }
+
@Provides
@ElasticsearchScope
public RateLimitedCache> writeCache(final HeroicReporter reporter) {
@@ -247,12 +262,21 @@ public SuggestBackend suggestBackend(Lazy kv) {
@Provides
@ElasticsearchScope
public LifeCycle life(
- LifeCycleManager manager, Lazy kv
+ LifeCycleManager manager, Lazy kv
) {
return manager.build(kv.get());
}
}
+ @Module
+ class O {
+ @Provides
+ @ElasticsearchScope
+ public Integer numSuggestionsLimit() {
+ return numSuggestionsLimit.getLimit();
+ }
+ }
+
@Override
public Optional id() {
return id;
@@ -275,6 +299,8 @@ public static class Builder {
private Optional templateName = empty();
private Optional backendType = empty();
private Optional configure = empty();
+ private Optional numSuggestionsLimit =
+ Optional.of(NumSuggestionsLimit.of());
public Builder id(final String id) {
checkNotNull(id, "id");
@@ -282,6 +308,12 @@ public Builder id(final String id) {
return this;
}
+ public Builder numSuggestionsLimit(final NumSuggestionsLimit numSuggestionsLimit) {
+ checkNotNull(numSuggestionsLimit, "numSuggestionsLimit");
+ this.numSuggestionsLimit = of(numSuggestionsLimit);
+ return this;
+ }
+
public Builder group(final Groups groups) {
checkNotNull(groups, "groups");
this.groups = of(groups);
@@ -347,19 +379,19 @@ public Builder configure(final boolean configure) {
public ElasticsearchSuggestModule build() {
return new ElasticsearchSuggestModule(
- id,
- groups,
- connection,
- writesPerSecond,
- rateLimitSlowStartSeconds,
- writeCacheDurationMinutes,
- writeCacheConcurrency,
- writeCacheMaxSize,
- distributedCacheSrvRecord,
- templateName,
- backendType,
- configure
- );
+ id,
+ groups,
+ connection,
+ writesPerSecond,
+ rateLimitSlowStartSeconds,
+ writeCacheDurationMinutes,
+ writeCacheConcurrency,
+ writeCacheMaxSize,
+ distributedCacheSrvRecord,
+ templateName,
+ backendType,
+ configure,
+ numSuggestionsLimit.get().asOptionalInt());
}
}
}
diff --git a/suggest/elasticsearch/src/main/java/com/spotify/heroic/suggest/elasticsearch/SuggestBackendKV.java b/suggest/elasticsearch/src/main/java/com/spotify/heroic/suggest/elasticsearch/SuggestBackendKV.java
index 235f3844f..32a0dbd66 100644
--- a/suggest/elasticsearch/src/main/java/com/spotify/heroic/suggest/elasticsearch/SuggestBackendKV.java
+++ b/suggest/elasticsearch/src/main/java/com/spotify/heroic/suggest/elasticsearch/SuggestBackendKV.java
@@ -58,8 +58,10 @@
import com.spotify.heroic.statistics.FutureReporter;
import com.spotify.heroic.statistics.SuggestBackendReporter;
import com.spotify.heroic.suggest.KeySuggest;
+import com.spotify.heroic.suggest.NumSuggestionsLimit;
import com.spotify.heroic.suggest.SuggestBackend;
import com.spotify.heroic.suggest.TagKeyCount;
+import com.spotify.heroic.suggest.TagKeyCount.Suggestion;
import com.spotify.heroic.suggest.TagSuggest;
import com.spotify.heroic.suggest.TagValueSuggest;
import com.spotify.heroic.suggest.TagValuesSuggest;
@@ -112,15 +114,18 @@
import org.elasticsearch.search.aggregations.BucketOrder;
import org.elasticsearch.search.aggregations.bucket.MultiBucketsAggregation;
import org.elasticsearch.search.aggregations.bucket.terms.Terms;
+import org.elasticsearch.search.aggregations.bucket.terms.Terms.Bucket;
import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
import org.elasticsearch.search.aggregations.metrics.Cardinality;
import org.elasticsearch.search.aggregations.metrics.CardinalityAggregationBuilder;
import org.elasticsearch.search.aggregations.metrics.TopHitsAggregationBuilder;
+import org.jetbrains.annotations.NotNull;
@ElasticsearchScope
public class SuggestBackendKV extends AbstractElasticsearchBackend
implements SuggestBackend, Grouped, LifeCycles {
+ private NumSuggestionsLimit numSuggestionsLimit = NumSuggestionsLimit.of();
private final Tracer tracer = Tracing.getTracer();
private static final String WRITE_CACHE_SIZE = "write-cache-size";
@@ -168,13 +173,15 @@ public SuggestBackendKV(
final SuggestBackendReporter reporter,
final RateLimitedCache> writeCache,
final Groups groups,
- @Named("configure") boolean configure) {
+ @Named("configure") boolean configure,
+ @Named("numSuggestionsLimit") Integer numSuggestionsLimit) {
super(async);
this.connection = connection;
this.reporter = reporter;
this.writeCache = writeCache;
this.groups = groups;
this.configure = configure;
+ this.numSuggestionsLimit = NumSuggestionsLimit.of(numSuggestionsLimit);
}
@Override
@@ -215,25 +222,16 @@ public AsyncFuture tagValuesSuggest(final TagValuesSuggest.Req
final QueryBuilder query = bool.hasClauses() ? bool : matchAllQuery();
SearchRequest searchRequest = c.getIndex().search(TAG_TYPE);
- searchRequest.source().size(0).query(query).timeout(TIMEOUT);
- final OptionalLimit limit = request.getLimit();
final OptionalLimit groupLimit = request.getGroupLimit();
- {
- final TermsAggregationBuilder terms =
- AggregationBuilders.terms("keys").field(TAG_SKEY_RAW);
- limit.asInteger().ifPresent(l -> terms.size(l + 1));
+ // use this.numSuggestionsLimit unless request.limit has been set.
+ final int numSuggestionsLimit =
+ this.numSuggestionsLimit.calculateNewLimit(request.getLimit());
- searchRequest.source().aggregation(terms);
- // make value bucket one entry larger than necessary to figure out when limiting
- // is applied.
- final TermsAggregationBuilder cardinality =
- AggregationBuilders.terms("values").field(TAG_SVAL_RAW);
+ searchRequest.source().size(numSuggestionsLimit).query(query).timeout(TIMEOUT);
- groupLimit.asInteger().ifPresent(l -> cardinality.size(l + 1));
- terms.subAggregation(cardinality);
- }
+ addCardinality(searchRequest, groupLimit);
final ResolvableFuture future = async.future();
@@ -241,38 +239,8 @@ public AsyncFuture tagValuesSuggest(final TagValuesSuggest.Req
return future.directTransform(
(SearchResponse response) -> {
- final List suggestions = new ArrayList<>();
-
- if (response.getAggregations() == null) {
- return new TagValuesSuggest(Collections.emptyList(), Boolean.FALSE);
- }
-
- final Terms terms = response.getAggregations().get("keys");
-
- // TODO: check type
- final List buckets = terms.getBuckets();
-
- for (final Terms.Bucket bucket : limit.limitList(buckets)) {
- final Terms valueTerms = bucket.getAggregations().get("values");
-
- // TODO: check type
- List valueBuckets = valueTerms.getBuckets();
-
- final SortedSet result = new TreeSet<>();
-
- for (final Terms.Bucket valueBucket : valueBuckets) {
- result.add(valueBucket.getKeyAsString());
- }
-
- final SortedSet values = groupLimit.limitSortedSet(result);
- final boolean limited = groupLimit.isGreater(valueBuckets.size());
-
- suggestions.add(new TagValuesSuggest.Suggestion(
- bucket.getKeyAsString(), values, limited));
- }
-
- return new TagValuesSuggest(
- ImmutableList.copyOf(suggestions), limit.isGreater(buckets.size()));
+ return createTagValuesSuggest(numSuggestionsLimit,
+ groupLimit, response);
});
});
}
@@ -281,60 +249,35 @@ public AsyncFuture tagValuesSuggest(final TagValuesSuggest.Req
public AsyncFuture tagValueSuggest(final TagValueSuggest.Request request) {
return connection.doto(
(final Connection c) -> {
- final BoolQueryBuilder bool = boolQuery();
+ final BoolQueryBuilder boolQueryBuilder = boolQuery();
final Optional key = request.getKey();
key.ifPresent(
k -> {
if (!k.isEmpty()) {
- bool.must(termQuery(TAG_SKEY_RAW, k));
+ boolQueryBuilder.must(termQuery(TAG_SKEY_RAW, k));
}
});
- QueryBuilder query = bool.hasClauses() ? bool : matchAllQuery();
-
- if (!(request.getFilter() instanceof TrueFilter)) {
- query = new BoolQueryBuilder().must(query).filter(filter(request.getFilter()));
- }
+ QueryBuilder query = augmentTagQueryBuilder(boolQueryBuilder, request.getFilter());
SearchRequest searchRequest = c.getIndex().search(TAG_TYPE);
- searchRequest.source().size(0).query(query);
- final OptionalLimit limit = request.getLimit();
+ final int suggestionsLimit = this.numSuggestionsLimit.calculateNewLimit(
+ request.getLimit());
+ searchRequest.source().size(suggestionsLimit).query(query);
- {
- final TermsAggregationBuilder terms =
- AggregationBuilders.terms("values")
- .field(TAG_SVAL_RAW)
- .order(BucketOrder.key(true));
-
- limit.asInteger().ifPresent(l -> terms.size(l + 1));
- searchRequest.source().aggregation(terms);
- }
+ aggregateRequestByTerms(searchRequest, suggestionsLimit);
final ResolvableFuture future = async.future();
c.execute(searchRequest, bind(future));
+ final var optionalNumSuggestionsLimit = this.numSuggestionsLimit.asOptionalLimit();
+
return future.directTransform(
(SearchResponse response) -> {
- final ImmutableList.Builder suggestions = ImmutableList.builder();
-
- if (response.getAggregations() == null) {
- return new TagValueSuggest(Collections.emptyList(), Boolean.FALSE);
- }
-
- final Terms terms = response.getAggregations().get("values");
-
- // TODO: Check type
- final List buckets = terms.getBuckets();
-
- for (final Terms.Bucket bucket : limit.limitList(buckets)) {
- suggestions.add(bucket.getKeyAsString());
- }
-
- return new TagValueSuggest(suggestions.build(),
- limit.isGreater(buckets.size()));
+ return createTagValueSuggest(optionalNumSuggestionsLimit, response);
});
});
}
@@ -346,34 +289,13 @@ public AsyncFuture tagKeyCount(final TagKeyCount.Request request) {
final QueryBuilder root = new BoolQueryBuilder().must(filter(request.getFilter()));
SearchRequest searchRequest = c.getIndex().search(TAG_TYPE);
+
searchRequest.source().size(0).query(root);
final OptionalLimit limit = request.getLimit();
final OptionalLimit exactLimit = request.getExactLimit();
- {
- final TermsAggregationBuilder keys =
- AggregationBuilders.terms("keys").field(TAG_SKEY_RAW);
- limit.asInteger().ifPresent(l -> keys.size(l + 1));
- searchRequest.source().aggregation(keys);
-
- final CardinalityAggregationBuilder cardinality =
- AggregationBuilders.cardinality("cardinality").field(TAG_SVAL_RAW);
-
- keys.subAggregation(cardinality);
-
- exactLimit
- .asInteger()
- .ifPresent(
- size -> {
- final TermsAggregationBuilder values = AggregationBuilders
- .terms("values")
- .field(TAG_SVAL_RAW)
- .size(size + 1);
-
- keys.subAggregation(values);
- });
- }
+ aggregateRequestByKeys(searchRequest, limit, exactLimit);
final ResolvableFuture future = async.future();
c.execute(searchRequest, bind(future));
@@ -391,26 +313,7 @@ public AsyncFuture tagKeyCount(final TagKeyCount.Request request) {
final List buckets = keys.getBuckets();
for (final Terms.Bucket bucket : limit.limitList(buckets)) {
- final Cardinality cardinality =
- bucket.getAggregations().get("cardinality");
- final Terms values = bucket.getAggregations().get("values");
-
- final Optional> exactValues;
-
- if (values != null) {
- exactValues =
- Optional.of(
- values.getBuckets().stream()
- .map(MultiBucketsAggregation.Bucket::getKeyAsString)
- .collect(Collectors.toSet()))
- .filter(sets -> !exactLimit.isGreater(sets.size()));
- } else {
- exactValues = Optional.empty();
- }
-
- suggestions.add(
- new TagKeyCount.Suggestion(
- bucket.getKeyAsString(), cardinality.getValue(), exactValues));
+ suggestions.add(createSuggestion(exactLimit, bucket));
}
return new TagKeyCount(
@@ -419,162 +322,37 @@ public AsyncFuture tagKeyCount(final TagKeyCount.Request request) {
});
}
- @Override
- public AsyncFuture tagSuggest(TagSuggest.Request request) {
- return connection.doto(
- (final Connection c) -> {
- final BoolQueryBuilder bool = boolQuery();
-
- final Optional key = request.getKey();
- final Optional value = request.getValue();
-
- // special case: key and value are equal, which indicates that _any_ match should be
- // in effect.
- // XXX: Enhance API to allow for this to be intentional instead of this by
- // introducing an 'any' field.
- if (key.isPresent() && value.isPresent() && key.equals(value)) {
- bool.should(matchTermKey(key.get()));
- bool.should(matchTermValue(value.get()));
- } else {
- key.ifPresent(
- k -> {
- if (!k.isEmpty()) {
- bool.must(matchTermKey(k).boost(2.0f));
- }
- });
-
- value.ifPresent(
- v -> {
- if (!v.isEmpty()) {
- bool.must(matchTermValue(v));
- }
- });
- }
-
- QueryBuilder query = bool.hasClauses() ? bool : matchAllQuery();
-
- if (!(request.getFilter() instanceof TrueFilter)) {
- query = new BoolQueryBuilder().must(query).filter(filter(request.getFilter()));
- }
-
- SearchRequest searchRequest = c.getIndex().search(TAG_TYPE);
- searchRequest.source().size(0).query(query).timeout(TIMEOUT);
-
- // aggregation
- {
- final TopHitsAggregationBuilder hits =
- AggregationBuilders.topHits("hits")
- .size(1)
- .fetchSource(TAG_SUGGEST_SOURCES, new String[0]);
-
- final TermsAggregationBuilder terms =
- AggregationBuilders.terms("terms").field(TAG_KV).subAggregation(hits);
-
- request.getLimit().asInteger().ifPresent(terms::size);
-
- searchRequest.source().aggregation(terms);
- }
-
- final ResolvableFuture future = async.future();
- c.execute(searchRequest, bind(future));
-
- return future.directTransform(
- (SearchResponse response) -> {
- final ImmutableList.Builder suggestions =
- ImmutableList.builder();
-
- final Aggregations aggregations = response.getAggregations();
-
- if (aggregations == null) {
- return new TagSuggest();
- }
-
- Stream> buckets =
- c.parseHits(aggregations.get("terms"));
-
- buckets.forEach(bucket -> {
- SearchHits hits = bucket.getSecond();
- final SearchHit hit = hits.getAt(0);
- final Map doc = hit.getSourceAsMap();
-
- final String k = (String) doc.get(TAG_SKEY);
- final String v = (String) doc.get(TAG_SVAL);
-
- suggestions.add(new TagSuggest.Suggestion(hits.getMaxScore(), k, v));
- });
-
- return new TagSuggest(suggestions.build());
- });
- });
- }
-
@Override
public AsyncFuture keySuggest(final KeySuggest.Request request) {
+ // Consider connection.doto to be an async version of using-with-resources
return connection.doto(
(final Connection c) -> {
BoolQueryBuilder bool = boolQuery();
- request
- .getKey()
- .ifPresent(
- k -> {
- if (!k.isEmpty()) {
- final String l = k.toLowerCase();
- final BoolQueryBuilder b = boolQuery();
- b.should(termQuery(SERIES_KEY_ANALYZED, l));
- b.should(termQuery(SERIES_KEY_PREFIX, l).boost(1.5f));
- b.should(termQuery(SERIES_KEY, l).boost(2.0f));
- bool.must(b);
- }
- });
+ addKeyToBuilder(request, bool);
- QueryBuilder query = bool.hasClauses() ? bool : matchAllQuery();
- if (!(request.getFilter() instanceof TrueFilter)) {
- query = new BoolQueryBuilder().must(query).filter(filter(request.getFilter()));
- }
+ QueryBuilder query = augmentTagQueryBuilder(bool, request.getFilter());
SearchRequest searchRequest = c.getIndex().search(SERIES_TYPE);
- searchRequest.source().size(0).query(query);
-
- // aggregation
- {
- final TopHitsAggregationBuilder hits =
- AggregationBuilders.topHits("hits")
- .size(1)
- .fetchSource(KEY_SUGGEST_SOURCES, new String[0]);
- final TermsAggregationBuilder keys =
- AggregationBuilders.terms("keys").field(KEY).subAggregation(hits);
+ final int suggestLimit = numSuggestionsLimit
+ .calculateNewLimit(request.getLimit());
+ searchRequest.source().size(suggestLimit).query(query);
- request.getLimit().asInteger().ifPresent(keys::size);
- searchRequest.source().aggregation(keys);
- }
+ // aggregation
+ addTermsToRequest(searchRequest, KEY_SUGGEST_SOURCES, "keys", KEY, suggestLimit);
final ResolvableFuture future = async.future();
c.execute(searchRequest, bind(future));
- return future.directTransform(
- (SearchResponse response) -> {
- final Set suggestions = new LinkedHashSet<>();
-
- if (response.getAggregations() == null) {
- return new KeySuggest(Collections.emptyList());
- }
-
- Stream> buckets =
- c.parseHits(response.getAggregations().get("keys"));
-
- buckets.forEach(bucket -> {
- SearchHits hits = bucket.getSecond();
- suggestions.add(new KeySuggest.Suggestion(
- hits.getMaxScore(), bucket.getFirst()));
- });
-
- return new KeySuggest(ImmutableList.copyOf(suggestions));
- });
+ return createKeySuggestAsync(c, future);
});
}
+ public NumSuggestionsLimit getNumSuggestionsLimit() {
+ return numSuggestionsLimit;
+ }
+
@Override
public AsyncFuture write(final WriteSuggest.Request request) {
return write(request, tracer.getCurrentSpan());
@@ -695,8 +473,8 @@ public AsyncFuture write(
ImmutableList.of(response.getTook().getMillis()),
ImmutableList.of());
})
- .onDone(writeContext)
- .onFinished(rootSpan::end);
+ .onDone(writeContext)
+ .onFinished(rootSpan::end);
});
}
@@ -705,6 +483,22 @@ public Statistics getStatistics() {
return new Statistics(WRITE_CACHE_SIZE, writeCache.size());
}
+ private static void addKeyToBuilder(KeySuggest.Request request, BoolQueryBuilder bool) {
+ request
+ .getKey()
+ .ifPresent(
+ k -> {
+ if (!k.isEmpty()) {
+ final String l = k.toLowerCase();
+ final BoolQueryBuilder b = boolQuery();
+ b.should(termQuery(SERIES_KEY_ANALYZED, l));
+ b.should(termQuery(SERIES_KEY_PREFIX, l).boost(1.5f));
+ b.should(termQuery(SERIES_KEY, l).boost(2.0f));
+ bool.must(b);
+ }
+ });
+ }
+
private AsyncFuture start() {
final AsyncFuture future = connection.start();
@@ -862,4 +656,321 @@ public static Supplier factory() {
public String toString() {
return "SuggestBackendKV(connection=" + this.connection + ")";
}
+
+ @NotNull
+ private static TagValuesSuggest createTagValuesSuggest(
+ int suggestValueLimit,
+ OptionalLimit groupLimit,
+ SearchResponse response) {
+
+ final List suggestions = new ArrayList<>();
+
+ if (response.getAggregations() == null) {
+ return new TagValuesSuggest(Collections.emptyList(), Boolean.FALSE);
+ }
+
+ final Terms terms = response.getAggregations().get("keys");
+
+ // TODO: check type
+ final List buckets = terms.getBuckets();
+
+ // WE need valueCount to be an object since
+ var suggestValueCount = new int[1];
+ for (final Terms.Bucket bucket : buckets) {
+ suggestions.add(
+ createTagValuesSuggestFromBuckets(bucket, suggestValueCount, groupLimit,
+ suggestValueLimit));
+ }
+
+ return new TagValuesSuggest(
+ ImmutableList.copyOf(suggestions), suggestValueLimit > buckets.size());
+ }
+
+ private static void addCardinality(SearchRequest searchRequest, OptionalLimit groupLimit) {
+ final TermsAggregationBuilder terms =
+ AggregationBuilders.terms("keys").field(TAG_SKEY_RAW);
+
+ searchRequest.source().aggregation(terms);
+
+ // make value bucket one entry larger than necessary to figure out when limiting
+ // is applied.
+ final TermsAggregationBuilder cardinality =
+ AggregationBuilders.terms("values").field(TAG_SVAL_RAW);
+
+ groupLimit.asInteger().ifPresent(l -> cardinality.size(l + 1));
+ terms.subAggregation(cardinality);
+ }
+
+ /**
+ * Creates a tagValuesSuggestion object.
+ *
+ * It respects both the groupLimit and the suggestValueCount limit. The latter
+ * increases/persists between calls to this method whereas the former just
+ * applies to this bucket.
+ *
+ * @param bucket used to get the Terms, from which we get value buckets
+ * @param suggestValueCount why on earth are we using an array of length 1 to
+ * track a count? because it must survive between calls
+ * and Integer objects aren't mutable. Also AtomicInt
+ * is very, very slow in a highly-multithreaded context.
+ * @param groupLimit only permit this many tag values for this bucket
+ * @param suggestValueLimit only permit this many tag values total, across
+ * all buckets.
+ * @return a suitably-constrained tag values suggestion object
+ */
+ private static TagValuesSuggest.Suggestion createTagValuesSuggestFromBuckets(
+ Bucket bucket, int[] suggestValueCount, OptionalLimit groupLimit,
+ int suggestValueLimit) {
+
+ final Terms valueTerms = bucket.getAggregations().get("values");
+
+ List valueBuckets = valueTerms.getBuckets();
+
+ SortedSet values = new TreeSet<>();
+
+ for (final Terms.Bucket valueBucket : valueBuckets) {
+ values.add(valueBucket.getKeyAsString());
+
+ suggestValueCount[0] = suggestValueCount[0] + 1;
+ if (suggestValueCount[0] == suggestValueLimit) {
+ break;
+ }
+ }
+
+ // any group of tag values may not exceed this limit
+ values = groupLimit.limitSortedSet(values);
+ final boolean limited = groupLimit.isGreater(valueBuckets.size()) ||
+ suggestValueCount[0] == suggestValueLimit;
+
+ return new TagValuesSuggest.Suggestion(
+ bucket.getKeyAsString(), values, limited);
+ }
+
+ private static AsyncFuture createKeySuggestAsync(Connection c,
+ ResolvableFuture future) {
+ return future.directTransform(
+ (SearchResponse response) -> {
+ final Set suggestions = new LinkedHashSet<>();
+
+ if (response.getAggregations() == null) {
+ return new KeySuggest(Collections.emptyList());
+ }
+
+ Stream> buckets =
+ c.parseHits(response.getAggregations().get("keys"));
+
+ buckets.forEach(bucket -> {
+ SearchHits hits = bucket.getSecond();
+ suggestions.add(new KeySuggest.Suggestion(
+ hits.getMaxScore(), bucket.getFirst()));
+ });
+
+ return new KeySuggest(ImmutableList.copyOf(suggestions));
+ });
+ }
+
+ @NotNull
+ private static TagValueSuggest createTagValueSuggest(OptionalLimit optionalNumSuggestionsLimit,
+ SearchResponse response) {
+ final ImmutableList.Builder suggestions = ImmutableList.builder();
+
+ if (response.getAggregations() == null) {
+ return new TagValueSuggest(Collections.emptyList(), Boolean.FALSE);
+ }
+
+ final Terms terms = response.getAggregations().get("values");
+
+ // TODO: Check type
+ final List buckets = terms.getBuckets();
+
+ for (final var bucket : optionalNumSuggestionsLimit.limitList(buckets)) {
+ suggestions.add(bucket.getKeyAsString());
+ }
+
+ return new TagValueSuggest(suggestions.build(),
+ optionalNumSuggestionsLimit.isGreater(buckets.size()));
+ }
+
+ private static void aggregateRequestByTerms(SearchRequest searchRequest,
+ int numSuggestionsLimit) {
+ final TermsAggregationBuilder terms =
+ AggregationBuilders.terms("values")
+ .size(numSuggestionsLimit)
+ .field(TAG_SVAL_RAW)
+ .order(BucketOrder.key(true));
+
+ searchRequest.source().aggregation(terms);
+ }
+
+ private static Suggestion createSuggestion(OptionalLimit exactLimit, Bucket bucket) {
+ final Cardinality cardinality =
+ bucket.getAggregations().get("cardinality");
+ final Terms values = bucket.getAggregations().get("values");
+
+ final Optional> exactValues;
+
+ if (values != null) {
+ exactValues =
+ Optional.of(
+ values.getBuckets().stream()
+ .map(MultiBucketsAggregation.Bucket::getKeyAsString)
+ .collect(Collectors.toSet()))
+ .filter(sets -> !exactLimit.isGreater(sets.size()));
+ } else {
+ exactValues = Optional.empty();
+ }
+
+ return new TagKeyCount.Suggestion(bucket.getKeyAsString(),
+ cardinality.getValue(), exactValues);
+ }
+
+ private static void aggregateRequestByKeys(
+ SearchRequest searchRequest, OptionalLimit limit, OptionalLimit exactLimit) {
+
+ final TermsAggregationBuilder keysAggBuilder = AggregationBuilders.terms("keys")
+ .field(TAG_SKEY_RAW);
+ limit.asInteger().ifPresent(l -> keysAggBuilder.size(l + 1));
+
+ // Assign keys aggregation to request
+ searchRequest.source().aggregation(keysAggBuilder);
+
+ final CardinalityAggregationBuilder cardinality =
+ AggregationBuilders.cardinality("cardinality").field(TAG_SVAL_RAW);
+
+ // Then add cardinality aggregation to the keys builder
+ keysAggBuilder.subAggregation(cardinality);
+
+ // Finally, if an exact limit is present, sub-aggregate by the supplied terms
+ exactLimit
+ .asInteger()
+ .ifPresent(
+ size -> {
+ final TermsAggregationBuilder values =
+ AggregationBuilders.terms("values").field(TAG_SVAL_RAW).size(size + 1);
+
+ keysAggBuilder.subAggregation(values);
+ });
+ }
+
+ @Override
+ public AsyncFuture tagSuggest(TagSuggest.Request request) {
+ return connection.doto(
+ (final Connection c) -> {
+ final BoolQueryBuilder boolQueryBuilder = boolQuery();
+
+ final Optional key = request.getKey();
+ final Optional value = request.getValue();
+
+ // special case: key and value are equal, which indicates that _any_ match should be
+ // in effect.
+ // XXX: Enhance API to allow for this to be intentional instead of this by
+ // introducing an 'any' field.
+ addKeyAndValueToBuilder(boolQueryBuilder, key, value);
+
+ var query = augmentTagQueryBuilder(boolQueryBuilder, request.getFilter());
+
+ SearchRequest searchRequest = c.getIndex().search(TAG_TYPE);
+
+ final int suggestionsLimit =
+ this.numSuggestionsLimit.calculateNewLimit(request.getLimit());
+
+ searchRequest
+ .source()
+ .size(suggestionsLimit)
+ .query(query)
+ .timeout(TIMEOUT);
+
+ // aggregation
+ addTermsToRequest(searchRequest, TAG_SUGGEST_SOURCES, "terms", TAG_KV,
+ suggestionsLimit);
+
+ final ResolvableFuture future = async.future();
+ c.execute(searchRequest, bind(future));
+
+ return getTagSuggestAsync(c, future);
+ });
+ }
+
+ private static AsyncFuture getTagSuggestAsync(Connection c,
+ ResolvableFuture future) {
+ return future.directTransform(
+ (SearchResponse response) -> {
+ final ImmutableList.Builder suggestions =
+ ImmutableList.builder();
+
+ final Aggregations aggregations = response.getAggregations();
+
+ if (aggregations == null) {
+ return new TagSuggest();
+ }
+
+ Stream> buckets =
+ c.parseHits(aggregations.get("terms"));
+
+ buckets.forEach(bucket -> {
+ SearchHits hits = bucket.getSecond();
+ final SearchHit hit = hits.getAt(0);
+ final Map doc = hit.getSourceAsMap();
+
+ final String k = (String) doc.get(TAG_SKEY);
+ final String v = (String) doc.get(TAG_SVAL);
+
+ suggestions.add(new TagSuggest.Suggestion(hits.getMaxScore(), k, v));
+ });
+
+ return new TagSuggest(suggestions.build());
+ });
+ }
+
+ private static void addTermsToRequest(SearchRequest searchRequest,
+ String[] tagSuggestSources, String term, String tagKv,
+ int suggestionsLimit) {
+ final TopHitsAggregationBuilder hits =
+ AggregationBuilders.topHits("hits")
+ .size(1)
+ .fetchSource(tagSuggestSources, new String[0]);
+
+ final TermsAggregationBuilder terms =
+ AggregationBuilders
+ .terms(term)
+ .size(suggestionsLimit)
+ .field(tagKv)
+ .subAggregation(hits);
+
+ searchRequest.source().aggregation(terms);
+ }
+
+ private static void addKeyAndValueToBuilder(BoolQueryBuilder bool, Optional key,
+ Optional value) {
+ if (key.isPresent() && value.isPresent() && key.equals(value)) {
+ bool.should(matchTermKey(key.get()));
+ bool.should(matchTermValue(value.get()));
+ } else {
+ key.ifPresent(
+ k -> {
+ if (!k.isEmpty()) {
+ bool.must(matchTermKey(k).boost(2.0f));
+ }
+ });
+
+ value.ifPresent(
+ v -> {
+ if (!v.isEmpty()) {
+ bool.must(matchTermValue(v));
+ }
+ });
+ }
+ }
+
+ private static QueryBuilder augmentTagQueryBuilder(BoolQueryBuilder boolQueryBuilder,
+ Filter filter) {
+ QueryBuilder query = boolQueryBuilder.hasClauses() ? boolQueryBuilder : matchAllQuery();
+
+ if (!(filter instanceof TrueFilter)) {
+ query = new BoolQueryBuilder().must(query).filter(filter(filter));
+ }
+
+ return query;
+ }
+
}
diff --git a/suggest/memory/src/main/java/com/spotify/heroic/suggest/memory/MemoryBackend.java b/suggest/memory/src/main/java/com/spotify/heroic/suggest/memory/MemoryBackend.java
index 29a10eec1..c05d259fb 100644
--- a/suggest/memory/src/main/java/com/spotify/heroic/suggest/memory/MemoryBackend.java
+++ b/suggest/memory/src/main/java/com/spotify/heroic/suggest/memory/MemoryBackend.java
@@ -30,6 +30,7 @@
import com.spotify.heroic.common.Series;
import com.spotify.heroic.filter.Filter;
import com.spotify.heroic.suggest.KeySuggest;
+import com.spotify.heroic.suggest.NumSuggestionsLimit;
import com.spotify.heroic.suggest.SuggestBackend;
import com.spotify.heroic.suggest.TagKeyCount;
import com.spotify.heroic.suggest.TagSuggest;
@@ -50,6 +51,7 @@
import java.util.SortedSet;
import java.util.TreeMap;
import java.util.TreeSet;
+import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;
@@ -57,6 +59,7 @@
import java.util.stream.Collectors;
import java.util.stream.Stream;
import javax.inject.Inject;
+import org.jetbrains.annotations.NotNull;
@MemoryScope
public class MemoryBackend implements SuggestBackend, Grouped {
@@ -75,10 +78,14 @@ public class MemoryBackend implements SuggestBackend, Grouped {
private final Groups groups;
private final AsyncFramework async;
+ private NumSuggestionsLimit numSuggestionsLimit = NumSuggestionsLimit.of();
+
@Inject
- public MemoryBackend(final Groups groups, final AsyncFramework async) {
+ public MemoryBackend(final Groups groups, final AsyncFramework async,
+ Integer numSuggestionsIntLimit) {
this.groups = groups;
this.async = async;
+ this.numSuggestionsLimit = NumSuggestionsLimit.of(numSuggestionsIntLimit);
}
@Override
@@ -88,35 +95,15 @@ public AsyncFuture configure() {
@Override
public AsyncFuture tagValuesSuggest(TagValuesSuggest.Request request) {
- final Map> counts = new HashMap<>();
-
- final OptionalLimit groupLimit = request.getGroupLimit();
-
- try (final Stream series = lookupSeries(request.getFilter())) {
- series.forEach(s -> {
- for (final Map.Entry e : s.getTags().entrySet()) {
- Set c = counts.get(e.getKey());
-
- if (c == null) {
- c = new HashSet<>();
- counts.put(e.getKey(), c);
- }
- if (groupLimit.isGreaterOrEqual(c.size())) {
- continue;
- }
+ final var tagsToValues = getTagsToValuesLimited(
+ request.getGroupLimit(), request.getLimit(), request.getFilter());
- c.add(e.getValue());
- }
- });
- }
-
- final List suggestions = ImmutableList.copyOf(request
- .getLimit()
- .limitStream(counts.entrySet().stream())
- .map(e -> new TagValuesSuggest.Suggestion(e.getKey(),
- ImmutableSortedSet.copyOf(e.getValue()), false))
- .iterator());
+ final var suggestions = ImmutableList.copyOf(
+ tagsToValues.entrySet().stream()
+ .map(e -> new TagValuesSuggest.Suggestion(e.getKey(),
+ ImmutableSortedSet.copyOf(e.getValue()), false))
+ .iterator());
return async.resolved(new TagValuesSuggest(suggestions, false));
}
@@ -164,23 +151,26 @@ public AsyncFuture tagSuggest(final TagSuggest.Request request) {
values.ifPresent(parts -> parts.forEach(
k -> ids.retainAll(tagValues.getOrDefault(k, ImmutableSet.of()))));
+ int limit = numSuggestionsLimit.calculateNewLimit(request.getLimit());
+
final List suggestions = ImmutableList.copyOf(
- ImmutableSortedSet.copyOf(request
- .getLimit()
- .limitStream(ids.stream())
- .map(tagIndex::get)
- .filter(Objects::nonNull)
- .map(d -> new TagSuggest.Suggestion(
- SCORE, d.getId().getKey(), d.getId().getValue()))
- .iterator()));
+ ImmutableSortedSet.copyOf(ids.stream()
+ .limit(limit)
+ .map(tagIndex::get)
+ .filter(Objects::nonNull)
+ .map(d -> new TagSuggest.Suggestion(
+ SCORE, d.getId().getKey(), d.getId().getValue()))
+ .iterator()));
return async.resolved(new TagSuggest(suggestions));
}
}
@Override
- public AsyncFuture keySuggest(final KeySuggest.Request request) {
- final Optional> analyzedKeys = request.getKey().map(MemoryBackend::analyze);
+ public AsyncFuture keySuggest(
+ final KeySuggest.Request request) {
+ final Optional> analyzedKeys =
+ request.getKey().map(MemoryBackend::analyze);
final Set ids;
@@ -191,26 +181,32 @@ public AsyncFuture keySuggest(final KeySuggest.Request request) {
k -> ids.retainAll(keys.getOrDefault(k, ImmutableSet.of()))));
}
- final List suggestions = ImmutableList.copyOf(request
- .getLimit()
- .limitStream(ids.stream())
- .map(d -> new KeySuggest.Suggestion(SCORE, d))
- .iterator());
+ int limit = numSuggestionsLimit.calculateNewLimit(request.getLimit());
+
+ final List suggestions =
+ ImmutableList.copyOf(ids.stream().limit(limit).map(
+ d -> new KeySuggest.Suggestion(SCORE, d)).iterator());
return async.resolved(new KeySuggest(suggestions));
}
@Override
- public AsyncFuture tagValueSuggest(final TagValueSuggest.Request request) {
+ public AsyncFuture tagValueSuggest(
+ final TagValueSuggest.Request request) {
try (final Stream docs = lookupTags(request.getFilter())) {
final Stream ids = docs.map(TagDocument::getId);
- final List values = request
- .getLimit()
- .limitStream(
- request.getKey().map(k -> ids.filter(id -> id.getKey().equals(k))).orElse(ids))
- .map(TagId::getValue)
- .collect(Collectors.toList());
+ int limit = numSuggestionsLimit.calculateNewLimit(request.getLimit());
+
+ final var tagIdStream = request
+ .getKey()
+ .map(k -> ids.filter(id -> id.getKey().equals(k)))
+ .orElse(ids);
+
+ final List values = tagIdStream
+ .limit(limit)
+ .map(TagId::getValue)
+ .collect(Collectors.toList());
return async.resolved(new TagValueSuggest(values, false));
}
@@ -322,8 +318,8 @@ private Stream lookupTags(final Filter filter) {
final Lock l = lock.readLock();
l.lock();
return tagIndex.values().stream()
- .filter(e -> filter.apply(e.getSeries()))
- .onClose(l::unlock);
+ .filter(e -> filter.apply(e.getSeries()))
+ .onClose(l::unlock);
}
private Stream lookupSeries(final Filter filter) {
@@ -332,6 +328,63 @@ private Stream lookupSeries(final Filter filter) {
return series.stream().filter(filter::apply).onClose(l::unlock);
}
+ /**
+ * @param groupLimit maximum number of groups allowed
+ * @param requestLimit total values limit, supplied by the request
+ * @param requestFilter defines which series will be returned
+ * @return { tag → { val1, val2 }, tag2 → { val2, val9 }, ...}
+ */
+ @NotNull
+ private Map> getTagsToValuesLimited(
+ OptionalLimit groupLimit, OptionalLimit requestLimit, Filter requestFilter) {
+
+ final Map> allTagsToValuesMap = new HashMap<>();
+
+ final int limit = numSuggestionsLimit.calculateNewLimit(requestLimit);
+
+ try (final Stream seriesStream = lookupSeries(requestFilter)) {
+ return populateLimitedTagsToValuesMap(groupLimit, limit, seriesStream);
+ }
+
+ }
+
+ private Map> populateLimitedTagsToValuesMap(
+ OptionalLimit groupLimit, int requestLimit, Stream seriesStream) {
+
+ // TODO find alternative as this is slow since it flushes processor core caches
+ // and synchronizes them
+ AtomicLong totalNumValues = new AtomicLong();
+ var allTagsToValuesMap = new HashMap>();
+
+ seriesStream.forEach(series -> {
+ for (final Map.Entry tagValuePair : series.getTags().entrySet()) {
+ Set values = allTagsToValuesMap.get(tagValuePair.getKey());
+
+ // If you've not seen this tag before, create a holder for its
+ // values
+ if (values == null) {
+ values = new HashSet<>();
+ allTagsToValuesMap.put(tagValuePair.getKey(), values);
+ }
+
+ final int numValues = values.size();
+
+ if (groupLimit.isGreaterOrEqual(numValues)) {
+ continue;
+ }
+
+ if (totalNumValues.incrementAndGet() > requestLimit) {
+ continue;
+ }
+
+ // Add the value to the collection of values for this tag
+ values.add(tagValuePair.getValue());
+ }
+ });
+
+ return allTagsToValuesMap;
+ }
+
public String toString() {
return "MemoryBackend()";
}
diff --git a/suggest/memory/src/main/java/com/spotify/heroic/suggest/memory/MemorySuggestModule.java b/suggest/memory/src/main/java/com/spotify/heroic/suggest/memory/MemorySuggestModule.java
index d05d8a0ee..94aa41550 100644
--- a/suggest/memory/src/main/java/com/spotify/heroic/suggest/memory/MemorySuggestModule.java
+++ b/suggest/memory/src/main/java/com/spotify/heroic/suggest/memory/MemorySuggestModule.java
@@ -31,6 +31,7 @@
import com.spotify.heroic.common.Groups;
import com.spotify.heroic.common.ModuleId;
import com.spotify.heroic.dagger.PrimaryComponent;
+import com.spotify.heroic.suggest.NumSuggestionsLimit;
import com.spotify.heroic.suggest.SuggestModule;
import dagger.Component;
import dagger.Module;
@@ -43,14 +44,17 @@ public final class MemorySuggestModule implements SuggestModule, DynamicModuleId
private final Optional id;
private final Groups groups;
+ private final NumSuggestionsLimit numSuggestionsLimit;
@JsonCreator
public MemorySuggestModule(
@JsonProperty("id") Optional id,
- @JsonProperty("groups") Optional groups
+ @JsonProperty("groups") Optional groups,
+ @JsonProperty("numSuggestionsIntLimit") Optional numSuggestionsIntLimit
) {
this.id = id;
this.groups = groups.orElseGet(Groups::empty).or(DEFAULT_GROUP);
+ this.numSuggestionsLimit = NumSuggestionsLimit.of(numSuggestionsIntLimit);
}
@Override
@@ -60,11 +64,12 @@ public Exposed module(PrimaryComponent primary, Depends depends, final String id
.primaryComponent(primary)
.depends(depends)
.m(new M())
+ .o(new O())
.build();
}
@MemoryScope
- @Component(modules = M.class,
+ @Component(modules = {M.class, O.class},
dependencies = {PrimaryComponent.class, Depends.class})
interface C extends Exposed {
@Override
@@ -85,13 +90,25 @@ public Optional id() {
return id;
}
+ @Module
+ class O {
+ @Provides
+ @MemoryScope
+ public Integer numSuggestionsLimit() {
+ return numSuggestionsLimit.getLimit();
+ }
+ }
+
public static Builder builder() {
return new Builder();
}
public static class Builder {
+
private Optional id = empty();
private Optional groups = empty();
+ private Optional numSuggestionsLimit =
+ Optional.of(NumSuggestionsLimit.of());
public Builder id(final String id) {
checkNotNull(id, "id");
@@ -99,6 +116,12 @@ public Builder id(final String id) {
return this;
}
+ public Builder numSuggestionsLimit(final NumSuggestionsLimit numSuggestionsLimit) {
+ checkNotNull(numSuggestionsLimit, "numSuggestionsLimit");
+ this.numSuggestionsLimit = of(numSuggestionsLimit);
+ return this;
+ }
+
public Builder group(final Groups groups) {
checkNotNull(groups, "groups");
this.groups = of(groups);
@@ -106,7 +129,8 @@ public Builder group(final Groups groups) {
}
public MemorySuggestModule build() {
- return new MemorySuggestModule(id, groups);
+ return new MemorySuggestModule(id, groups,
+ numSuggestionsLimit.get().asOptionalInt());
}
}
}