Skip to content

Commit a306b9a

Browse files
chirinoandreaTP
chirino
and
andreaTP
authored
Directory backed runtime compiler cache (#1077)
This expands on #1041 and adds a built in Directory Cache implementation. When you use it with the runtime compiler: ``` var cache = new DirectoryCache(Path.of("/cache")); var instance = Instance.builder(module) .withMachineFactory( MachineFactoryCompiler.builder(module).withCache(cache).compile()) .build(); ``` It will create the files like: ``` /cache/.tmp/f-9979505440111212447.tmp ``` and then move the `/cache/.tmp/f-9979505440111212447.tmp` file in an atomic operation to `/cache/sha256/kr/gytkcm43c34ksqta8gmddw4ycfquc2g0qfifcpb-w.jar` (the file path is based on the sha256 of the original wasm file). It should be safe to use the same cache location by different processes. If any error occurs while generating the code into the temp dir, the temp file will get deleted. But if the process is killed before we can clean up, the temp file will not get cleaned up. User may want manually clean the /cache/.tmp directory when no process is writing to the cache. we use 2 character prefix of the sha to avoid hitting directory scaling issues with file systems. So if you have a VERY large cache all the cache entries will be partitioned over 4096 dirs. --------- Signed-off-by: Hiram Chirino <hiram@hiramchirino.com> Co-authored-by: andreatp <andrea.peruffo1982@gmail.com>
1 parent dfb93c5 commit a306b9a

13 files changed

Lines changed: 776 additions & 13 deletions

File tree

‎bom/pom.xml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,11 @@
4949
<artifactId>compiler</artifactId>
5050
<version>${project.version}</version>
5151
</dependency>
52+
<dependency>
53+
<groupId>com.dylibso.chicory</groupId>
54+
<artifactId>dircache-experimental</artifactId>
55+
<version>${project.version}</version>
56+
</dependency>
5257
<dependency>
5358
<groupId>com.dylibso.chicory</groupId>
5459
<artifactId>log</artifactId>

‎compiler/src/main/java/com/dylibso/chicory/compiler/Cache.java‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
package com.dylibso.chicory.compiler;
2+
3+
import java.io.IOException;
4+
5+
public interface Cache {
6+
7+
/**
8+
* Return the cached data for the given key if it exists else null.
9+
*
10+
* @param key "algo:digest"
11+
*/
12+
byte[] get(String key) throws IOException;
13+
14+
/**
15+
* Atomically publish data into the cache location for the key.
16+
* If another thread/process already published for this key then this is a no-op.
17+
*
18+
* @param key "algo:digest"
19+
* @param data the data to cache
20+
*/
21+
void putIfAbsent(String key, byte[] data) throws IOException;
22+
}

‎compiler/src/main/java/com/dylibso/chicory/compiler/MachineFactoryCompiler.java‎

Lines changed: 121 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,20 @@
44
import com.dylibso.chicory.compiler.internal.MachineFactory;
55
import com.dylibso.chicory.runtime.Instance;
66
import com.dylibso.chicory.runtime.Machine;
7+
import com.dylibso.chicory.wasm.ChicoryException;
78
import com.dylibso.chicory.wasm.WasmModule;
9+
import java.io.ByteArrayInputStream;
10+
import java.io.ByteArrayOutputStream;
11+
import java.io.IOException;
12+
import java.io.UncheckedIOException;
13+
import java.util.HashMap;
14+
import java.util.Map;
15+
import java.util.Properties;
816
import java.util.Set;
917
import java.util.function.Function;
18+
import java.util.jar.JarEntry;
19+
import java.util.jar.JarInputStream;
20+
import java.util.jar.JarOutputStream;
1021

1122
/**
1223
* Compiles WASM function bodies to JVM byte code that can be used as a machine factory for {@link Instance}'s.
@@ -38,7 +49,6 @@ public static Machine compile(Instance instance) {
3849
* Compiles a machine factory that can used in instance builders.
3950
* The module is only compiled once and the machine factory is reused for every
4051
* instance created by the builder.
41-
* <p>
4252
* <pre>
4353
* var module = Parser.parse(is);
4454
* var builder = Instance.builder(module)
@@ -56,7 +66,6 @@ public static Function<Instance, Machine> compile(WasmModule module) {
5666
* The builder allows you to configure the compiler options used to compile the module to
5767
* byte code.
5868
* This should be used when you want to create multiple instances of the same module.
59-
* <p>
6069
* <pre>
6170
* var module = Parser.parse(is);
6271
* var builder = Instance.builder(module)
@@ -76,6 +85,7 @@ public static Builder builder(WasmModule module) {
7685
public static final class Builder {
7786
private final WasmModule module;
7887
private final com.dylibso.chicory.compiler.internal.Compiler.Builder compilerBuilder;
88+
private Cache cache;
7989

8090
private Builder(WasmModule module) {
8191
this.module = module;
@@ -102,14 +112,116 @@ public Builder withInterpretedFunctions(Set<Integer> interpretedFunctions) {
102112
return this;
103113
}
104114

115+
public Builder withCache(Cache cache) {
116+
this.cache = cache;
117+
return this;
118+
}
119+
105120
public Function<Instance, Machine> compile() {
106-
var result =
107-
compilerBuilder
108-
.withClassCollectorFactory(ClassLoadingCollector::new)
109-
.build()
110-
.compile();
111-
var collector = (ClassLoadingCollector) result.collector();
112-
return new MachineFactory(module, collector.machineFactory());
121+
try {
122+
123+
// Can we load the byte codes from the cache?
124+
var useCache = cache != null && module.digest() != null;
125+
if (useCache) {
126+
byte[] cachedData = cache.get(module.digest());
127+
if (cachedData != null) {
128+
var collector = loadClassLoadingCollector(cachedData);
129+
return new MachineFactory(module, collector.machineFactory());
130+
}
131+
}
132+
133+
// Compile the byte codes...
134+
var result =
135+
compilerBuilder
136+
.withClassCollectorFactory(ClassLoadingCollector::new)
137+
.build()
138+
.compile();
139+
var collector = (ClassLoadingCollector) result.collector();
140+
141+
if (useCache) {
142+
// store results in the cache to speed the next time.
143+
cache.putIfAbsent(module.digest(), storeClassLoadingCollector(collector));
144+
}
145+
146+
return new MachineFactory(module, collector.machineFactory());
147+
} catch (IOException e) {
148+
throw new ChicoryException(e);
149+
}
150+
}
151+
}
152+
153+
private static byte[] storeClassLoadingCollector(ClassLoadingCollector collector) {
154+
try {
155+
// Create JAR in memory
156+
ByteArrayOutputStream baos = new ByteArrayOutputStream();
157+
try (JarOutputStream jos = new JarOutputStream(baos)) {
158+
// Store the properties file
159+
var properties = new Properties();
160+
properties.put("mainClass", collector.mainClassName());
161+
ByteArrayOutputStream propsBaos = new ByteArrayOutputStream();
162+
properties.store(propsBaos, "");
163+
164+
JarEntry propsEntry = new JarEntry("wasm-module.properties");
165+
jos.putNextEntry(propsEntry);
166+
jos.write(propsBaos.toByteArray());
167+
jos.closeEntry();
168+
169+
// Store all class files
170+
for (var entry : collector.classBytes().entrySet()) {
171+
var className = entry.getKey().replace('.', '/') + ".class";
172+
JarEntry classEntry = new JarEntry(className);
173+
jos.putNextEntry(classEntry);
174+
jos.write(entry.getValue());
175+
jos.closeEntry();
176+
}
177+
}
178+
return baos.toByteArray();
179+
} catch (IOException e) {
180+
throw new UncheckedIOException(e);
181+
}
182+
}
183+
184+
private static ClassLoadingCollector loadClassLoadingCollector(byte[] jarData)
185+
throws IOException {
186+
var collector = new ClassLoadingCollector();
187+
188+
// It was previously compiled, just load it from JAR.
189+
try (JarInputStream jis = new JarInputStream(new ByteArrayInputStream(jarData))) {
190+
var properties = new Properties();
191+
String mainClass = null;
192+
Map<String, byte[]> classes = new HashMap<>();
193+
194+
JarEntry entry;
195+
while ((entry = jis.getNextJarEntry()) != null) {
196+
if (entry.getName().equals("wasm-module.properties")) {
197+
// Load properties
198+
properties.load(jis);
199+
mainClass = properties.getProperty("mainClass");
200+
} else if (entry.getName().endsWith(".class")) {
201+
// Load class file
202+
String className = entry.getName().replace('/', '.').replace(".class", "");
203+
204+
// Read class file bytes
205+
ByteArrayOutputStream baos = new ByteArrayOutputStream();
206+
byte[] buffer = new byte[4096];
207+
int bytesRead;
208+
while ((bytesRead = jis.read(buffer)) != -1) {
209+
baos.write(buffer, 0, bytesRead);
210+
}
211+
classes.put(className, baos.toByteArray());
212+
}
213+
jis.closeEntry();
214+
}
215+
216+
// Add classes to collector
217+
for (var classEntry : classes.entrySet()) {
218+
if (classEntry.getKey().equals(mainClass)) {
219+
collector.putMainClass(mainClass, classEntry.getValue());
220+
} else {
221+
collector.put(classEntry.getKey(), classEntry.getValue());
222+
}
223+
}
113224
}
225+
return collector;
114226
}
115227
}

‎compiler/src/test/java/com/dylibso/chicory/compiler/internal/CacheTest.java‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
package com.dylibso.chicory.compiler.internal;
2+
3+
import static java.nio.charset.StandardCharsets.UTF_8;
4+
import static org.junit.jupiter.api.Assertions.assertEquals;
5+
6+
import com.dylibso.chicory.compiler.Cache;
7+
import com.dylibso.chicory.compiler.MachineFactoryCompiler;
8+
import com.dylibso.chicory.runtime.Instance;
9+
import com.dylibso.chicory.wasm.Parser;
10+
import java.io.IOException;
11+
import java.util.concurrent.ConcurrentHashMap;
12+
import java.util.concurrent.atomic.AtomicInteger;
13+
import org.junit.jupiter.api.Test;
14+
15+
public class CacheTest {
16+
17+
static class MockCache implements Cache {
18+
ConcurrentHashMap<String, byte[]> cache = new ConcurrentHashMap<>();
19+
20+
@Override
21+
public byte[] get(String key) throws IOException {
22+
return cache.get(key);
23+
}
24+
25+
@Override
26+
public void putIfAbsent(String key, byte[] data) throws IOException {
27+
cache.putIfAbsent(key, data);
28+
}
29+
}
30+
31+
public static class CacheWithHitCounter implements Cache {
32+
private final Cache cache;
33+
public AtomicInteger hits = new AtomicInteger(0);
34+
35+
public CacheWithHitCounter(Cache cache) {
36+
this.cache = cache;
37+
}
38+
39+
@Override
40+
public byte[] get(String key) throws IOException {
41+
byte[] result = cache.get(key);
42+
if (result != null) {
43+
hits.incrementAndGet();
44+
}
45+
return result;
46+
}
47+
48+
@Override
49+
public void putIfAbsent(String key, byte[] data) throws IOException {
50+
cache.putIfAbsent(key, data);
51+
}
52+
}
53+
54+
private void exerciseCountVowels(Instance instance) {
55+
var alloc = instance.export("alloc");
56+
var dealloc = instance.export("dealloc");
57+
var countVowels = instance.export("count_vowels");
58+
var memory = instance.memory();
59+
var message = "Hello, World!";
60+
var len = message.getBytes(UTF_8).length;
61+
int ptr = (int) alloc.apply(len)[0];
62+
memory.writeString(ptr, message);
63+
var result = countVowels.apply(ptr, len);
64+
dealloc.apply(ptr, len);
65+
assertEquals(3L, result[0]);
66+
}
67+
68+
@Test
69+
public void shouldCacheCompiledResultInMem() {
70+
71+
var cache = new CacheWithHitCounter(new MockCache());
72+
var module =
73+
Parser.parse(CacheTest.class.getResourceAsStream("/compiled/count_vowels.rs.wasm"));
74+
75+
var instance1 =
76+
Instance.builder(module)
77+
.withMachineFactory(
78+
MachineFactoryCompiler.builder(module).withCache(cache).compile())
79+
.build();
80+
81+
exerciseCountVowels(instance1);
82+
assertEquals(0, cache.hits.get());
83+
var instance2 =
84+
Instance.builder(module)
85+
.withMachineFactory(
86+
MachineFactoryCompiler.builder(module).withCache(cache).compile())
87+
.build();
88+
exerciseCountVowels(instance2);
89+
assertEquals(1, cache.hits.get());
90+
}
91+
}

‎dircache/pom.xml‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
<?xml version="1.0" encoding="UTF-8"?>
2+
<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
3+
<modelVersion>4.0.0</modelVersion>
4+
5+
<parent>
6+
<groupId>com.dylibso.chicory</groupId>
7+
<artifactId>chicory</artifactId>
8+
<version>999-SNAPSHOT</version>
9+
</parent>
10+
<artifactId>dircache-experimental</artifactId>
11+
<packaging>jar</packaging>
12+
13+
<name>Chicory - Dircache - Experimental</name>
14+
<description>Chicory directory based Cache</description>
15+
16+
<dependencies>
17+
<dependency>
18+
<groupId>com.dylibso.chicory</groupId>
19+
<artifactId>compiler</artifactId>
20+
</dependency>
21+
<dependency>
22+
<groupId>com.dylibso.chicory</groupId>
23+
<artifactId>runtime</artifactId>
24+
<scope>test</scope>
25+
</dependency>
26+
<dependency>
27+
<groupId>com.dylibso.chicory</groupId>
28+
<artifactId>wasm</artifactId>
29+
<scope>test</scope>
30+
</dependency>
31+
<dependency>
32+
<groupId>com.dylibso.chicory</groupId>
33+
<artifactId>wasm-corpus</artifactId>
34+
<scope>test</scope>
35+
</dependency>
36+
<dependency>
37+
<groupId>io.roastedroot</groupId>
38+
<artifactId>zerofs</artifactId>
39+
<scope>test</scope>
40+
</dependency>
41+
<dependency>
42+
<groupId>org.junit.jupiter</groupId>
43+
<artifactId>junit-jupiter-api</artifactId>
44+
<scope>test</scope>
45+
</dependency>
46+
<dependency>
47+
<groupId>org.junit.jupiter</groupId>
48+
<artifactId>junit-jupiter-engine</artifactId>
49+
<scope>test</scope>
50+
</dependency>
51+
52+
</dependencies>
53+
54+
</project>

0 commit comments

Comments
 (0)