Skip to content

Commit 28b99f0

Browse files
davidkyledavidkyle
authored
[ML] JIindex: Limit the size of bulk migrations (#36481)
1 parent ed26ae8 commit 28b99f0

File tree

2 files changed

+154
-8
lines changed

2 files changed

+154
-8
lines changed

‎x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/MlConfigMigrator.java‎

Lines changed: 74 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -38,11 +38,13 @@
3838
import java.util.Collection;
3939
import java.util.HashMap;
4040
import java.util.HashSet;
41+
import java.util.Iterator;
4142
import java.util.List;
4243
import java.util.Map;
4344
import java.util.Set;
4445
import java.util.concurrent.atomic.AtomicBoolean;
4546
import java.util.concurrent.atomic.AtomicReference;
47+
import java.util.function.Function;
4648
import java.util.stream.Collectors;
4749

4850
import static org.elasticsearch.xpack.core.ClientHelper.ML_ORIGIN;
@@ -68,7 +70,10 @@
6870
* If there was an error in step 3 and the config is in both the clusterstate and
6971
* index then when the migrator retries it must not overwrite an existing job config
7072
* document as once the index document is present all update operations will function
71-
* on that rather than the clusterstate
73+
* on that rather than the clusterstate.
74+
*
75+
* The number of configs indexed in each bulk operation is limited by {@link #MAX_BULK_WRITE_SIZE}
76+
* pairs of datafeeds and jobs are migrated together.
7277
*/
7378
public class MlConfigMigrator {
7479

@@ -77,6 +82,8 @@ public class MlConfigMigrator {
7782
public static final String MIGRATED_FROM_VERSION = "migrated from version";
7883
public static final Version MIN_NODE_VERSION = Version.V_6_6_0;
7984

85+
static final int MAX_BULK_WRITE_SIZE = 100;
86+
8087
private final Client client;
8188
private final ClusterService clusterService;
8289

@@ -118,12 +125,15 @@ public void migrateConfigsWithoutTasks(ClusterState clusterState, ActionListener
118125
return;
119126
}
120127

128+
121129
logger.debug("migrating ml configurations");
122130

123-
Collection<DatafeedConfig> datafeedsToMigrate = stoppedDatafeedConfigs(clusterState);
124-
List<Job> jobsToMigrate = nonDeletingJobs(closedJobConfigs(clusterState)).stream()
131+
Collection<DatafeedConfig> stoppedDatafeeds = stoppedDatafeedConfigs(clusterState);
132+
Map<String, Job> eligibleJobs = nonDeletingJobs(closedJobConfigs(clusterState)).stream()
125133
.map(MlConfigMigrator::updateJobForMigration)
126-
.collect(Collectors.toList());
134+
.collect(Collectors.toMap(Job::getId, Function.identity(), (a, b) -> a));
135+
136+
JobsAndDatafeeds jobsAndDatafeedsToMigrate = limitWrites(stoppedDatafeeds, eligibleJobs);
127137

128138
ActionListener<Boolean> unMarkMigrationInProgress = ActionListener.wrap(
129139
response -> {
@@ -136,16 +146,16 @@ public void migrateConfigsWithoutTasks(ClusterState clusterState, ActionListener
136146
}
137147
);
138148

139-
if (datafeedsToMigrate.isEmpty() && jobsToMigrate.isEmpty()) {
149+
if (jobsAndDatafeedsToMigrate.totalCount() == 0) {
140150
unMarkMigrationInProgress.onResponse(Boolean.FALSE);
141151
return;
142152
}
143153

144-
writeConfigToIndex(datafeedsToMigrate, jobsToMigrate, ActionListener.wrap(
154+
writeConfigToIndex(jobsAndDatafeedsToMigrate.datafeedConfigs, jobsAndDatafeedsToMigrate.jobs, ActionListener.wrap(
145155
failedDocumentIds -> {
146-
List<String> successfulJobWrites = filterFailedJobConfigWrites(failedDocumentIds, jobsToMigrate);
156+
List<String> successfulJobWrites = filterFailedJobConfigWrites(failedDocumentIds, jobsAndDatafeedsToMigrate.jobs);
147157
List<String> successfulDatafeedWrites =
148-
filterFailedDatafeedConfigWrites(failedDocumentIds, datafeedsToMigrate);
158+
filterFailedDatafeedConfigWrites(failedDocumentIds, jobsAndDatafeedsToMigrate.datafeedConfigs);
149159
removeFromClusterState(successfulJobWrites, successfulDatafeedWrites, unMarkMigrationInProgress);
150160
},
151161
unMarkMigrationInProgress::onFailure
@@ -350,6 +360,62 @@ public static List<DatafeedConfig> stoppedDatafeedConfigs(ClusterState clusterSt
350360
.collect(Collectors.toList());
351361
}
352362

363+
public static class JobsAndDatafeeds {
364+
List<Job> jobs;
365+
List<DatafeedConfig> datafeedConfigs;
366+
367+
private JobsAndDatafeeds() {
368+
jobs = new ArrayList<>();
369+
datafeedConfigs = new ArrayList<>();
370+
}
371+
372+
public int totalCount() {
373+
return jobs.size() + datafeedConfigs.size();
374+
}
375+
}
376+
377+
/**
378+
* Return at most {@link #MAX_BULK_WRITE_SIZE} configs favouring
379+
* datafeed and job pairs so if a datafeed is chosen so is its job.
380+
*
381+
* @param datafeedsToMigrate Datafeed configs
382+
* @param jobsToMigrate Job configs
383+
* @return Job and datafeed configs
384+
*/
385+
public static JobsAndDatafeeds limitWrites(Collection<DatafeedConfig> datafeedsToMigrate, Map<String, Job> jobsToMigrate) {
386+
JobsAndDatafeeds jobsAndDatafeeds = new JobsAndDatafeeds();
387+
388+
if (datafeedsToMigrate.size() + jobsToMigrate.size() <= MAX_BULK_WRITE_SIZE) {
389+
jobsAndDatafeeds.jobs.addAll(jobsToMigrate.values());
390+
jobsAndDatafeeds.datafeedConfigs.addAll(datafeedsToMigrate);
391+
return jobsAndDatafeeds;
392+
}
393+
394+
int count = 0;
395+
396+
// prioritise datafeed and job pairs
397+
for (DatafeedConfig datafeedConfig : datafeedsToMigrate) {
398+
if (count < MAX_BULK_WRITE_SIZE) {
399+
jobsAndDatafeeds.datafeedConfigs.add(datafeedConfig);
400+
count++;
401+
Job datafeedsJob = jobsToMigrate.remove(datafeedConfig.getJobId());
402+
if (datafeedsJob != null) {
403+
jobsAndDatafeeds.jobs.add(datafeedsJob);
404+
count++;
405+
}
406+
}
407+
}
408+
409+
// are there jobs without datafeeds to migrate
410+
Iterator<Job> iter = jobsToMigrate.values().iterator();
411+
while (iter.hasNext() && count < MAX_BULK_WRITE_SIZE) {
412+
jobsAndDatafeeds.jobs.add(iter.next());
413+
count++;
414+
}
415+
416+
return jobsAndDatafeeds;
417+
}
418+
353419
/**
354420
* Check for failures in the bulk response and return the
355421
* Ids of any documents not written to the index

‎x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/MlConfigMigratorTests.java‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,11 @@
2828
import java.util.ArrayList;
2929
import java.util.Arrays;
3030
import java.util.Collections;
31+
import java.util.HashMap;
3132
import java.util.List;
33+
import java.util.Map;
3234
import java.util.Set;
35+
import java.util.stream.Collectors;
3336

3437
import static org.hamcrest.Matchers.contains;
3538
import static org.hamcrest.Matchers.containsInAnyOrder;
@@ -333,6 +336,83 @@ public void testDatafeedIsEligibleForMigration_givenStoppedDatafeed() {
333336
assertTrue(MlConfigMigrator.datafeedIsEligibleForMigration(datafeedId, clusterState));
334337
}
335338

339+
public void testLimitWrites_GivenBelowLimit() {
340+
MlConfigMigrator.JobsAndDatafeeds jobsAndDatafeeds = MlConfigMigrator.limitWrites(Collections.emptyList(), Collections.emptyMap());
341+
assertThat(jobsAndDatafeeds.datafeedConfigs, empty());
342+
assertThat(jobsAndDatafeeds.jobs, empty());
343+
344+
List<DatafeedConfig> datafeeds = new ArrayList<>();
345+
Map<String, Job> jobs = new HashMap<>();
346+
347+
int numDatafeeds = MlConfigMigrator.MAX_BULK_WRITE_SIZE / 2;
348+
for (int i=0; i<numDatafeeds; i++) {
349+
String jobId = "job" + i;
350+
jobs.put(jobId, JobTests.buildJobBuilder(jobId).build());
351+
datafeeds.add(createCompatibleDatafeed(jobId));
352+
}
353+
354+
jobsAndDatafeeds = MlConfigMigrator.limitWrites(datafeeds, jobs);
355+
assertThat(jobsAndDatafeeds.datafeedConfigs, hasSize(numDatafeeds));
356+
assertThat(jobsAndDatafeeds.jobs, hasSize(numDatafeeds));
357+
}
358+
359+
public void testLimitWrites_GivenAboveLimit() {
360+
List<DatafeedConfig> datafeeds = new ArrayList<>();
361+
Map<String, Job> jobs = new HashMap<>();
362+
363+
int numDatafeeds = MlConfigMigrator.MAX_BULK_WRITE_SIZE / 2 + 10;
364+
for (int i=0; i<numDatafeeds; i++) {
365+
String jobId = "job" + i;
366+
jobs.put(jobId, JobTests.buildJobBuilder(jobId).build());
367+
datafeeds.add(createCompatibleDatafeed(jobId));
368+
}
369+
370+
MlConfigMigrator.JobsAndDatafeeds jobsAndDatafeeds = MlConfigMigrator.limitWrites(datafeeds, jobs);
371+
assertEquals(MlConfigMigrator.MAX_BULK_WRITE_SIZE, jobsAndDatafeeds.totalCount());
372+
assertThat(jobsAndDatafeeds.datafeedConfigs, hasSize(MlConfigMigrator.MAX_BULK_WRITE_SIZE / 2));
373+
assertThat(jobsAndDatafeeds.jobs, hasSize(MlConfigMigrator.MAX_BULK_WRITE_SIZE / 2));
374+
375+
// assert that for each datafeed its corresponding job is selected
376+
Set<String> selectedJobIds = jobsAndDatafeeds.jobs.stream().map(Job::getId).collect(Collectors.toSet());
377+
Set<String> datafeedJobIds = jobsAndDatafeeds.datafeedConfigs.stream().map(DatafeedConfig::getJobId).collect(Collectors.toSet());
378+
assertEquals(selectedJobIds, datafeedJobIds);
379+
}
380+
381+
public void testLimitWrites_GivenMoreJobsThanDatafeeds() {
382+
List<DatafeedConfig> datafeeds = new ArrayList<>();
383+
Map<String, Job> jobs = new HashMap<>();
384+
385+
int numDatafeeds = MlConfigMigrator.MAX_BULK_WRITE_SIZE / 2 - 10;
386+
for (int i=0; i<numDatafeeds; i++) {
387+
String jobId = "job" + i;
388+
jobs.put(jobId, JobTests.buildJobBuilder(jobId).build());
389+
datafeeds.add(createCompatibleDatafeed(jobId));
390+
}
391+
392+
for (int i=numDatafeeds; i<numDatafeeds + 40; i++) {
393+
String jobId = "job" + i;
394+
jobs.put(jobId, JobTests.buildJobBuilder(jobId).build());
395+
}
396+
397+
MlConfigMigrator.JobsAndDatafeeds jobsAndDatafeeds = MlConfigMigrator.limitWrites(datafeeds, jobs);
398+
assertEquals(MlConfigMigrator.MAX_BULK_WRITE_SIZE, jobsAndDatafeeds.totalCount());
399+
assertThat(jobsAndDatafeeds.datafeedConfigs, hasSize(numDatafeeds));
400+
assertThat(jobsAndDatafeeds.jobs, hasSize(MlConfigMigrator.MAX_BULK_WRITE_SIZE - numDatafeeds));
401+
402+
// assert that for each datafeed its corresponding job is selected
403+
Set<String> selectedJobIds = jobsAndDatafeeds.jobs.stream().map(Job::getId).collect(Collectors.toSet());
404+
Set<String> datafeedJobIds = jobsAndDatafeeds.datafeedConfigs.stream().map(DatafeedConfig::getJobId).collect(Collectors.toSet());
405+
assertTrue(selectedJobIds.containsAll(datafeedJobIds));
406+
}
407+
408+
public void testLimitWrites_GivenNullJob() {
409+
List<DatafeedConfig> datafeeds = Collections.singletonList(createCompatibleDatafeed("no-job-for-this-datafeed"));
410+
MlConfigMigrator.JobsAndDatafeeds jobsAndDatafeeds = MlConfigMigrator.limitWrites(datafeeds, Collections.emptyMap());
411+
412+
assertThat(jobsAndDatafeeds.datafeedConfigs, hasSize(1));
413+
assertThat(jobsAndDatafeeds.jobs, empty());
414+
}
415+
336416
private DatafeedConfig createCompatibleDatafeed(String jobId) {
337417
// create a datafeed without aggregations or anything
338418
// else that may cause validation errors

0 commit comments

Comments
 (0)