org.openrewrite.maven.cache.RocksdbMavenPomCache Maven / Gradle / Ivy
Show all versions of rewrite-maven Show documentation
/*
* Copyright 2020 the original author or authors.
*
* Licensed 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
*
* https://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 org.openrewrite.maven.cache;
import com.fasterxml.jackson.annotation.JsonAutoDetect;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.SerializationFeature;
import com.fasterxml.jackson.databind.cfg.ConstructorDetector;
import com.fasterxml.jackson.databind.json.JsonMapper;
import com.fasterxml.jackson.dataformat.smile.SmileFactory;
import com.fasterxml.jackson.dataformat.smile.SmileGenerator;
import com.fasterxml.jackson.datatype.jdk8.Jdk8Module;
import com.fasterxml.jackson.module.paramnames.ParameterNamesModule;
import org.jspecify.annotations.Nullable;
import org.openrewrite.maven.MavenDownloadingException;
import org.openrewrite.maven.tree.*;
import org.rocksdb.*;
import java.io.File;
import java.io.IOException;
import java.io.UncheckedIOException;
import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
/**
* Implementation of the maven cache that leverages Rocksdb. The keys and values are serialized to/from byte arrays
* using jackson. Things to know about this cache implementation:
*
*
It will create a rocks db in the workspace directory passed to it.
* If two caches are pointed to the same workspace folder, they will "share" the same underlying rocks database,
* it is thread-safe.
* Because multiple caches can share the same database, the close on this cache implementation does nothing.
* The database is closed via a system shutdown hook registered by this class. Any unexpected process termination
* is non-fatal, any non-flushed data is lost, but the database will not be corrupted.
* The database is configured to auto-flush when the in-memory size reaches 1MB.
* Rocksdb's write ahead log has been disabled because we are using this as a cache and do not need to recover any
* "lost" data.
* Rocksdb computes checksums for all of its files, normally it checks those on startup, this has been disabled as
* well.
*/
@SuppressWarnings("OptionalAssignedToNull")
public class RocksdbMavenPomCache implements MavenPomCache {
private static final String MODEL_VERSION_KEY = "org.openrewrite.maven.internal.Pom.version";
static ObjectMapper mapper;
//The RocksDB instance is thread safe, the first call to create a database for a workspace will open the database
//subsequent calls will get the same instances back. This cache also registers a shutdown hook to close the
//databases on shutdown.
private static final Map cacheMap = new HashMap<>();
static {
SmileFactory f = new SmileFactory();
f.configure(SmileGenerator.Feature.CHECK_SHARED_STRING_VALUES, true);
ObjectMapper m = JsonMapper.builder(f)
.constructorDetector(ConstructorDetector.USE_PROPERTIES_BASED)
.build()
.registerModule(new ParameterNamesModule())
.registerModule(new Jdk8Module())
.disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES)
.disable(SerializationFeature.FAIL_ON_EMPTY_BEANS)
.setSerializationInclusion(JsonInclude.Include.NON_NULL);
mapper = m.setVisibility(m.getSerializationConfig().getDefaultVisibilityChecker()
.withFieldVisibility(JsonAutoDetect.Visibility.ANY)
.withGetterVisibility(JsonAutoDetect.Visibility.NONE)
.withSetterVisibility(JsonAutoDetect.Visibility.NONE)
.withCreatorVisibility(JsonAutoDetect.Visibility.PUBLIC_ONLY));
//Init the rockdb native jni library
RocksDB.loadLibrary();
}
static synchronized RocksCache getCache(String pomCacheDir) {
return cacheMap.computeIfAbsent(pomCacheDir, RocksCache::new);
}
static synchronized void closeCache(String pomCacheDir) {
RocksCache cache = cacheMap.remove(pomCacheDir);
if (cache != null) {
cache.close();
}
}
private final RocksCache cache;
public RocksdbMavenPomCache(Path workspace) {
File pomCacheDir = new File(workspace.toFile(), ".rewrite-cache");
if (!pomCacheDir.exists() && !pomCacheDir.mkdirs()) {
throw new IllegalStateException("Unable to find or create maven pom cache at " + pomCacheDir);
} else if (!pomCacheDir.isDirectory()) {
throw new IllegalStateException("The maven pom cache workspace must be a directory at " + pomCacheDir);
}
// In case a stale lock file is left over from a previous run that was interrupted
File lock = new File(pomCacheDir, "LOCK");
if (lock.exists()) {
//noinspection ResultOfMethodCallIgnored
lock.delete();
}
cache = getCache(pomCacheDir.getAbsolutePath());
}
@Override
public @Nullable ResolvedPom getResolvedDependencyPom(ResolvedGroupArtifactVersion dependency) {
return null;
}
@Override
public void putResolvedDependencyPom(ResolvedGroupArtifactVersion dependency, ResolvedPom resolved) {
}
@Override
public @Nullable Optional getMavenMetadata(URI repo, GroupArtifactVersion gav) {
//The Maven metadata is not something that should be stored long term, as it will change over time.
return null;
}
@Override
public void putMavenMetadata(URI repo, GroupArtifactVersion gav, @Nullable MavenMetadata metadata) {
//The Maven metadata is not something that should be stored long term, as it will change over time.
}
@Override
public Optional getPom(ResolvedGroupArtifactVersion gav) throws MavenDownloadingException {
try {
return deserializePom(cache.get(serialize(gav.toString().getBytes(StandardCharsets.UTF_8))));
} catch (RocksDBException e) {
throw new MavenDownloadingException("Failed to deserialize POM from RocksDB cache", e,
new GroupArtifactVersion(gav.getGroupId(), gav.getArtifactId(), gav.getVersion()));
}
}
@Override
public void putPom(ResolvedGroupArtifactVersion gav, Pom pom) {
if (pom == null) {
return;
}
try {
cache.put(serialize(gav.toString().getBytes(StandardCharsets.UTF_8)), serialize(pom));
} catch (RocksDBException e) {
throw new IllegalStateException("Failed to save POM into RocksDB cache", e);
}
}
@Override
public @Nullable Optional getNormalizedRepository(MavenRepository repository) {
return null;
}
@Override
public void putNormalizedRepository(MavenRepository repository, MavenRepository normalized) {
}
static byte[] serialize(T object) {
if (object == null) {
return null;
}
try {
return mapper.writeValueAsBytes(object);
} catch (JsonProcessingException e) {
throw new IllegalArgumentException("Unable to serialize object to byte array.");
}
}
static @Nullable Optional deserializePom(byte[] bytes) {
try {
return bytes == null ? null : Optional.of(mapper.readValue(bytes, Pom.class));
} catch (IOException e) {
throw new UncheckedIOException(e);
}
}
/**
* Wrapper class around the rocksdb. The database and options are all backed by C++ data structures that
* must be explicitly closed to ensure proper memory management. Note, if the same database is being used
* by multiple threads, close should only be executed once all threads are done using the database. This class
* registers a shutdown hook to close all open databases, so there should be no need to explicitly close the
* databases.
*/
static class RocksCache {
private RocksDB database;
private final Options options;
private final WriteOptions writeOptions;
private final String cacheFolder;
RocksCache(String pomCacheDir) {
try {
options = new Options();
options.setCreateIfMissing(true);
//Default memtable buffer size is 64MB, changing this to 1MB because we are only caching pom.xml files
//When the memtable exceeds 1MB, rocks will write the contents to disk. Note, closing the database
//also forces a flush to occur.
options.setWriteBufferSize(1_000_000);
//since we are only using rocks db as a cache, turning off checksum verification when opening.
options.setParanoidChecks(false);
options.setParanoidFileChecks(false);
//Turn off write ahead log, there is no needs to record the data in both memory and in a log (from which
//rocks can recover in the case of a system failure).
writeOptions = new WriteOptions();
writeOptions.setDisableWAL(true);
database = RocksDB.open(options, pomCacheDir);
cacheFolder = pomCacheDir;
} catch (RocksDBException exception) {
throw new IllegalStateException("Unable to create cache database." + exception.getMessage(), exception);
}
try {
cleanCacheIfCorrupt(pomCacheDir);
} catch (Exception ex) {
throw new IllegalStateException("Unable to clear corrupt maven pom cache.", ex);
}
}
private void cleanCacheIfCorrupt(String pomCacheDir) {
boolean clearCache = false;
try {
database.verifyChecksum();
} catch (RocksDBException ex) {
clearCache = true;
}
if (!clearCache) {
// Validate the model version used by the cache matches the version of RawMaven
clearCache = !modelVersionMatches();
}
if (clearCache) {
database.close();
try {
//Delete the contents of the database.
RocksDB.destroyDB(pomCacheDir, options);
} catch (RocksDBException exception) {
throw new IllegalStateException("Unable to clear maven pom cache at " + pomCacheDir, exception);
}
// Re-open the database anew
try {
database = RocksDB.open(options, pomCacheDir);
} catch (RocksDBException exception) {
throw new IllegalStateException("Unable to clear maven pom cache at " + pomCacheDir, exception);
}
}
// Update the model version key
updateDatabaseModelVersion(Pom.getModelVersion());
}
private boolean modelVersionMatches() {
// Check if there are any keys in the database, there is no need to check the model version if
// this is a new database.
RocksIterator iterator = database.newIterator();
iterator.seekToFirst();
try {
if (iterator.isValid()) {
// Check the model version used by the cache against the in-memory version
try {
byte[] rawVersion = database.get(serialize(MODEL_VERSION_KEY));
Integer databaseVersion = rawVersion == null ? null : mapper.readValue(rawVersion, Integer.class);
if (databaseVersion == null || Pom.getModelVersion() != databaseVersion) {
return false;
}
} catch (IOException exception) {
throw new IllegalStateException("Unable to deserialize database model version.", exception);
}
}
} catch (RocksDBException rocksException) {
throw new IllegalStateException("Unable verify database model version.", rocksException);
}
return true;
}
protected void updateDatabaseModelVersion(Integer newVersion) {
try {
put(serialize(MODEL_VERSION_KEY), serialize(newVersion));
} catch (RocksDBException e) {
throw new IllegalStateException("Unable to update database model version.");
}
}
private void put(byte[] key, byte[] value) throws RocksDBException {
database.put(writeOptions, key, value);
}
private byte[] get(byte[] key) throws RocksDBException {
return database.get(key);
}
private void close() {
// This will flush any in-memory memtables to disk and free up resources held
// by the underlying C++ code. The worse case scenario is that this is not called because the system exits
// abnormally, in which case, the data in-memory is simply not saved to the cache.
if (Files.exists(Paths.get(cacheFolder))) {
//Attempting to close the database if the file has been deleted underneath it prevents the process
//from exiting.
try(FlushOptions flushOptions = new FlushOptions()) {
database.flush(flushOptions);
} catch (RocksDBException e) {
//Error while flushing the database (before close), this means that any interim changes in memory
//will not be persisted. This is non-fatal.
} finally {
database.close();
}
}
writeOptions.close();
options.close();
}
}
}