Skip to content

Commit 41a47c2

Browse files
samxbrsamxbr
authored
Make GeoIP database node loader project-aware (#129829)
- loads the downloaded GeoIP databases from system index to ingest node file system for each project - each project's databases are loaded to directory `tmp/geoip-databases/{nodeId}/{projectId}`
1 parent 73d2e30 commit 41a47c2

17 files changed

+497
-305
lines changed

‎modules/ingest-geoip/src/internalClusterTest/java/org/elasticsearch/ingest/geoip/DatabaseNodeServiceIT.java

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,9 @@
1515
import org.elasticsearch.action.admin.indices.flush.FlushRequest;
1616
import org.elasticsearch.action.admin.indices.refresh.RefreshRequest;
1717
import org.elasticsearch.action.index.IndexRequest;
18+
import org.elasticsearch.cluster.metadata.ProjectId;
1819
import org.elasticsearch.common.hash.MessageDigests;
20+
import org.elasticsearch.core.FixForMultiProject;
1921
import org.elasticsearch.xcontent.XContentType;
2022

2123
import java.io.BufferedInputStream;
@@ -40,6 +42,10 @@
4042
import static org.hamcrest.Matchers.equalTo;
4143

4244
public class DatabaseNodeServiceIT extends AbstractGeoIpIT {
45+
46+
@FixForMultiProject(description = "Use random project ID after ESIntegTestCase is MP enabled")
47+
private final ProjectId projectId = ProjectId.DEFAULT;
48+
4349
/*
4450
* This test makes sure that if we index an ordinary mmdb file into the .geoip_databases index, it is correctly handled upon retrieval.
4551
*/
@@ -50,15 +56,15 @@ public void testNonGzippedDatabase() throws Exception {
5056
String databaseName = randomAlphaOfLength(20) + "-" + databaseFileName;
5157
byte[] mmdbBytes = getBytesForFile(databaseFileName);
5258
final DatabaseNodeService databaseNodeService = internalCluster().getInstance(DatabaseNodeService.class);
53-
assertNull(databaseNodeService.getDatabase(databaseName));
59+
assertNull(databaseNodeService.getDatabase(projectId, databaseName));
5460
int numChunks = indexData(databaseName, mmdbBytes);
5561
/*
5662
* If DatabaseNodeService::checkDatabases runs it will sometimes (rarely) remove the database we are using in this test while we
5763
* are trying to assert things about it. So if it does then we 'just' try again.
5864
*/
5965
assertBusy(() -> {
6066
retrieveDatabase(databaseNodeService, databaseName, mmdbBytes, numChunks);
61-
assertNotNull(databaseNodeService.getDatabase(databaseName));
67+
assertNotNull(databaseNodeService.getDatabase(projectId, databaseName));
6268
assertValidDatabase(databaseNodeService, databaseName, databaseType);
6369
});
6470
}
@@ -75,15 +81,15 @@ public void testGzippedDatabase() throws Exception {
7581
byte[] mmdbBytes = getBytesForFile(databaseFileName);
7682
byte[] gzipBytes = gzipFileBytes(databaseName, mmdbBytes);
7783
final DatabaseNodeService databaseNodeService = internalCluster().getInstance(DatabaseNodeService.class);
78-
assertNull(databaseNodeService.getDatabase(databaseName));
84+
assertNull(databaseNodeService.getDatabase(projectId, databaseName));
7985
int numChunks = indexData(databaseName, gzipBytes);
8086
/*
8187
* If DatabaseNodeService::checkDatabases runs it will sometimes (rarely) remove the database we are using in this test while we
8288
* are trying to assert things about it. So if it does then we 'just' try again.
8389
*/
8490
assertBusy(() -> {
8591
retrieveDatabase(databaseNodeService, databaseName, gzipBytes, numChunks);
86-
assertNotNull(databaseNodeService.getDatabase(databaseName));
92+
assertNotNull(databaseNodeService.getDatabase(projectId, databaseName));
8793
assertValidDatabase(databaseNodeService, databaseName, databaseType);
8894
});
8995
}
@@ -93,7 +99,7 @@ public void testGzippedDatabase() throws Exception {
9399
*/
94100
private void assertValidDatabase(DatabaseNodeService databaseNodeService, String databaseFileName, String databaseType)
95101
throws IOException {
96-
IpDatabase database = databaseNodeService.getDatabase(databaseFileName);
102+
IpDatabase database = databaseNodeService.getDatabase(projectId, databaseFileName);
97103
assertNotNull(database);
98104
assertThat(database.getDatabaseType(), equalTo(databaseType));
99105
CountryResponse countryResponse = database.getResponse("89.160.20.128", GeoIpTestUtils::getCountry);
@@ -110,7 +116,7 @@ private void assertValidDatabase(DatabaseNodeService databaseNodeService, String
110116
private void retrieveDatabase(DatabaseNodeService databaseNodeService, String databaseFileName, byte[] expectedBytes, int numChunks)
111117
throws IOException {
112118
GeoIpTaskState.Metadata metadata = new GeoIpTaskState.Metadata(1, 0, numChunks - 1, getMd5(expectedBytes), 1);
113-
databaseNodeService.retrieveAndUpdateDatabase(databaseFileName, metadata);
119+
databaseNodeService.retrieveAndUpdateDatabase(projectId, databaseFileName, metadata);
114120
}
115121

116122
private String getMd5(byte[] bytes) {

‎modules/ingest-geoip/src/internalClusterTest/java/org/elasticsearch/ingest/geoip/ReloadingDatabasesWhilePerformingGeoLookupsIT.java

Lines changed: 55 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,11 @@
1111

1212
import org.apache.lucene.tests.util.LuceneTestCase;
1313
import org.elasticsearch.client.internal.Client;
14-
import org.elasticsearch.cluster.ClusterState;
14+
import org.elasticsearch.cluster.ClusterName;
15+
import org.elasticsearch.cluster.metadata.ProjectId;
16+
import org.elasticsearch.cluster.metadata.ProjectMetadata;
17+
import org.elasticsearch.cluster.project.ProjectResolver;
18+
import org.elasticsearch.cluster.project.TestProjectResolvers;
1519
import org.elasticsearch.cluster.service.ClusterService;
1620
import org.elasticsearch.common.util.concurrent.AtomicArray;
1721
import org.elasticsearch.core.IOUtils;
@@ -32,6 +36,7 @@
3236
import java.util.concurrent.atomic.AtomicInteger;
3337
import java.util.concurrent.atomic.AtomicReference;
3438

39+
import static org.elasticsearch.cluster.ClusterState.builder;
3540
import static org.elasticsearch.ingest.geoip.GeoIpProcessor.GEOIP_TYPE;
3641
import static org.elasticsearch.ingest.geoip.GeoIpTestUtils.copyDatabase;
3742
import static org.elasticsearch.ingest.geoip.GeoIpTestUtils.copyDefaultDatabases;
@@ -62,31 +67,39 @@ public class ReloadingDatabasesWhilePerformingGeoLookupsIT extends ESTestCase {
6267
* geoip processor instance is using the related {@link DatabaseReaderLazyLoader} instance
6368
*/
6469
public void test() throws Exception {
70+
ProjectId projectId = randomProjectIdOrDefault();
6571
Path geoIpConfigDir = createTempDir();
6672
Path geoIpTmpDir = createTempDir();
6773
ClusterService clusterService = mock(ClusterService.class);
68-
when(clusterService.state()).thenReturn(ClusterState.EMPTY_STATE);
69-
DatabaseNodeService databaseNodeService = createRegistry(geoIpConfigDir, geoIpTmpDir, clusterService);
74+
when(clusterService.state()).thenReturn(
75+
builder(ClusterName.DEFAULT).putProjectMetadata(ProjectMetadata.builder(projectId).build()).build()
76+
);
77+
DatabaseNodeService databaseNodeService = createRegistry(
78+
geoIpConfigDir,
79+
geoIpTmpDir,
80+
clusterService,
81+
TestProjectResolvers.singleProject(projectId)
82+
);
7083
GeoIpProcessor.Factory factory = new GeoIpProcessor.Factory(GEOIP_TYPE, databaseNodeService);
7184
copyDatabase("GeoLite2-City-Test.mmdb", geoIpTmpDir.resolve("GeoLite2-City.mmdb"));
7285
copyDatabase("GeoLite2-City-Test.mmdb", geoIpTmpDir.resolve("GeoLite2-City-Test.mmdb"));
73-
databaseNodeService.updateDatabase("GeoLite2-City.mmdb", "md5", geoIpTmpDir.resolve("GeoLite2-City.mmdb"));
74-
databaseNodeService.updateDatabase("GeoLite2-City-Test.mmdb", "md5", geoIpTmpDir.resolve("GeoLite2-City-Test.mmdb"));
75-
lazyLoadReaders(databaseNodeService);
86+
databaseNodeService.updateDatabase(projectId, "GeoLite2-City.mmdb", "md5", geoIpTmpDir.resolve("GeoLite2-City.mmdb"));
87+
databaseNodeService.updateDatabase(projectId, "GeoLite2-City-Test.mmdb", "md5", geoIpTmpDir.resolve("GeoLite2-City-Test.mmdb"));
88+
lazyLoadReaders(projectId, databaseNodeService);
7689

7790
final GeoIpProcessor processor1 = (GeoIpProcessor) factory.create(
7891
null,
7992
"_tag",
8093
null,
8194
new HashMap<>(Map.of("field", "_field")),
82-
null
95+
projectId
8396
);
8497
final GeoIpProcessor processor2 = (GeoIpProcessor) factory.create(
8598
null,
8699
"_tag",
87100
null,
88101
new HashMap<>(Map.of("field", "_field", "database_file", "GeoLite2-City-Test.mmdb")),
89-
null
102+
projectId
90103
);
91104

92105
final AtomicBoolean completed = new AtomicBoolean(false);
@@ -134,9 +147,9 @@ public void test() throws Exception {
134147
Thread updateDatabaseThread = new Thread(() -> {
135148
for (int i = 0; i < numberOfDatabaseUpdates; i++) {
136149
try {
137-
DatabaseReaderLazyLoader previous1 = databaseNodeService.get("GeoLite2-City.mmdb");
150+
DatabaseReaderLazyLoader previous1 = databaseNodeService.get(projectId, "GeoLite2-City.mmdb");
138151
if (Files.exists(geoIpTmpDir.resolve("GeoLite2-City.mmdb"))) {
139-
databaseNodeService.removeStaleEntries(List.of("GeoLite2-City.mmdb"));
152+
databaseNodeService.removeStaleEntries(projectId, List.of("GeoLite2-City.mmdb"));
140153
assertBusy(() -> {
141154
// lazy loader may still be in use by an ingest thread,
142155
// wait for any potential ingest thread to release the lazy loader (DatabaseReaderLazyLoader#postLookup(...)),
@@ -146,22 +159,32 @@ public void test() throws Exception {
146159
});
147160
} else {
148161
copyDatabase("GeoLite2-City-Test.mmdb", geoIpTmpDir.resolve("GeoLite2-City.mmdb"));
149-
databaseNodeService.updateDatabase("GeoLite2-City.mmdb", "md5", geoIpTmpDir.resolve("GeoLite2-City.mmdb"));
162+
databaseNodeService.updateDatabase(
163+
projectId,
164+
"GeoLite2-City.mmdb",
165+
"md5",
166+
geoIpTmpDir.resolve("GeoLite2-City.mmdb")
167+
);
150168
}
151-
DatabaseReaderLazyLoader previous2 = databaseNodeService.get("GeoLite2-City-Test.mmdb");
169+
DatabaseReaderLazyLoader previous2 = databaseNodeService.get(projectId, "GeoLite2-City-Test.mmdb");
152170
copyDatabase(
153171
i % 2 == 0 ? "GeoIP2-City-Test.mmdb" : "GeoLite2-City-Test.mmdb",
154172
geoIpTmpDir.resolve("GeoLite2-City-Test.mmdb")
155173
);
156-
databaseNodeService.updateDatabase("GeoLite2-City-Test.mmdb", "md5", geoIpTmpDir.resolve("GeoLite2-City-Test.mmdb"));
174+
databaseNodeService.updateDatabase(
175+
projectId,
176+
"GeoLite2-City-Test.mmdb",
177+
"md5",
178+
geoIpTmpDir.resolve("GeoLite2-City-Test.mmdb")
179+
);
157180

158-
DatabaseReaderLazyLoader current1 = databaseNodeService.get("GeoLite2-City.mmdb");
159-
DatabaseReaderLazyLoader current2 = databaseNodeService.get("GeoLite2-City-Test.mmdb");
181+
DatabaseReaderLazyLoader current1 = databaseNodeService.get(projectId, "GeoLite2-City.mmdb");
182+
DatabaseReaderLazyLoader current2 = databaseNodeService.get(projectId, "GeoLite2-City-Test.mmdb");
160183
assertThat(current1, not(sameInstance(previous1)));
161184
assertThat(current2, not(sameInstance(previous2)));
162185

163186
// lazy load type and reader:
164-
lazyLoadReaders(databaseNodeService);
187+
lazyLoadReaders(projectId, databaseNodeService);
165188
} catch (Exception | AssertionError e) {
166189
logger.error("error in update databases thread after run [" + i + "]", e);
167190
failureHolder2.set(e);
@@ -193,8 +216,12 @@ public void test() throws Exception {
193216
IOUtils.rm(geoIpConfigDir, geoIpTmpDir);
194217
}
195218

196-
private static DatabaseNodeService createRegistry(Path geoIpConfigDir, Path geoIpTmpDir, ClusterService clusterService)
197-
throws IOException {
219+
private static DatabaseNodeService createRegistry(
220+
Path geoIpConfigDir,
221+
Path geoIpTmpDir,
222+
ClusterService clusterService,
223+
ProjectResolver projectResolver
224+
) throws IOException {
198225
GeoIpCache cache = new GeoIpCache(0);
199226
ConfigDatabases configDatabases = new ConfigDatabases(geoIpConfigDir, cache);
200227
copyDefaultDatabases(geoIpConfigDir, configDatabases);
@@ -204,19 +231,21 @@ private static DatabaseNodeService createRegistry(Path geoIpConfigDir, Path geoI
204231
cache,
205232
configDatabases,
206233
Runnable::run,
207-
clusterService
234+
clusterService,
235+
mock(IngestService.class),
236+
projectResolver
208237
);
209-
databaseNodeService.initialize("nodeId", mock(ResourceWatcherService.class), mock(IngestService.class));
238+
databaseNodeService.initialize("nodeId", mock(ResourceWatcherService.class));
210239
return databaseNodeService;
211240
}
212241

213-
private static void lazyLoadReaders(DatabaseNodeService databaseNodeService) throws IOException {
214-
if (databaseNodeService.get("GeoLite2-City.mmdb") != null) {
215-
databaseNodeService.get("GeoLite2-City.mmdb").getDatabaseType();
216-
databaseNodeService.get("GeoLite2-City.mmdb").getResponse("2.125.160.216", GeoIpTestUtils::getCity);
242+
private static void lazyLoadReaders(ProjectId projectId, DatabaseNodeService databaseNodeService) throws IOException {
243+
if (databaseNodeService.get(projectId, "GeoLite2-City.mmdb") != null) {
244+
databaseNodeService.get(projectId, "GeoLite2-City.mmdb").getDatabaseType();
245+
databaseNodeService.get(projectId, "GeoLite2-City.mmdb").getResponse("2.125.160.216", GeoIpTestUtils::getCity);
217246
}
218-
databaseNodeService.get("GeoLite2-City-Test.mmdb").getDatabaseType();
219-
databaseNodeService.get("GeoLite2-City-Test.mmdb").getResponse("2.125.160.216", GeoIpTestUtils::getCity);
247+
databaseNodeService.get(projectId, "GeoLite2-City-Test.mmdb").getDatabaseType();
248+
databaseNodeService.get(projectId, "GeoLite2-City-Test.mmdb").getResponse("2.125.160.216", GeoIpTestUtils::getCity);
220249
}
221250

222251
}

‎modules/ingest-geoip/src/main/java/org/elasticsearch/ingest/geoip/ConfigDatabases.java

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@
1010

1111
import org.apache.logging.log4j.LogManager;
1212
import org.apache.logging.log4j.Logger;
13+
import org.elasticsearch.cluster.metadata.ProjectId;
14+
import org.elasticsearch.core.FixForMultiProject;
1315
import org.elasticsearch.env.Environment;
1416
import org.elasticsearch.watcher.FileChangesListener;
1517
import org.elasticsearch.watcher.FileWatcher;
@@ -69,12 +71,13 @@ Map<String, DatabaseReaderLazyLoader> getConfigDatabases() {
6971
return configDatabases;
7072
}
7173

74+
@FixForMultiProject(description = "Replace DEFAULT project")
7275
void updateDatabase(Path file, boolean update) {
7376
String databaseFileName = file.getFileName().toString();
7477
try {
7578
if (update) {
7679
logger.info("database file changed [{}], reloading database...", file);
77-
DatabaseReaderLazyLoader loader = new DatabaseReaderLazyLoader(cache, file, null);
80+
DatabaseReaderLazyLoader loader = new DatabaseReaderLazyLoader(ProjectId.DEFAULT, cache, file, null);
7881
DatabaseReaderLazyLoader existing = configDatabases.put(databaseFileName, loader);
7982
if (existing != null) {
8083
existing.shutdown();
@@ -90,6 +93,7 @@ void updateDatabase(Path file, boolean update) {
9093
}
9194
}
9295

96+
@FixForMultiProject(description = "Replace DEFAULT project")
9397
Map<String, DatabaseReaderLazyLoader> initConfigDatabases() throws IOException {
9498
Map<String, DatabaseReaderLazyLoader> databases = new HashMap<>();
9599

@@ -103,7 +107,7 @@ Map<String, DatabaseReaderLazyLoader> initConfigDatabases() throws IOException {
103107
if (Files.isRegularFile(databasePath) && pathMatcher.matches(databasePath)) {
104108
assert Files.exists(databasePath);
105109
String databaseFileName = databasePath.getFileName().toString();
106-
DatabaseReaderLazyLoader loader = new DatabaseReaderLazyLoader(cache, databasePath, null);
110+
DatabaseReaderLazyLoader loader = new DatabaseReaderLazyLoader(ProjectId.DEFAULT, cache, databasePath, null);
107111
databases.put(databaseFileName, loader);
108112
}
109113
}

0 commit comments

Comments
 (0)