ValkeyCacheService: pluggable ObjectInputStream for cross-relocation reads

pom.xml

@@ -16,7 +16,7 @@
 	<artifactId>objectify</artifactId>
 	<!-- Streak fork: upstream master (PR #527, Valkey cache) stamped as a fixed
 	     version for the mailfoogae Artifact Registry. Bump -streak-valkey-N for hotfixes. -->
-	<version>6.1.4-streak-valkey-7</version>
+	<version>6.1.4-streak-valkey-8</version>
 
 	<name>Objectify App Engine</name>
 	<description>The simplest convenient interface to the Google App Engine datastore</description>

src/main/java/com/googlecode/objectify/cache/valkey/ValkeyCacheService.java

@@ -15,6 +15,7 @@
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.io.InputStream;
 import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
 import java.nio.charset.StandardCharsets;
@@ -25,6 +26,7 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
@@ -80,6 +82,19 @@
  */
 public class ValkeyCacheService implements MemcacheService {
 
+	/**
+	 * Wraps the raw bytes of a stored value in the {@link ObjectInputStream} used to deserialize it.
+	 * The default is a plain {@link ObjectInputStream}, but callers can supply a subclass that
+	 * overrides {@link ObjectInputStream#resolveClass} — Streak passes one that bridges shaded
+	 * ("streak."-relocated) and un-relocated class names, so processes built with different shadowJar
+	 * relocations can share one cache. Kept as an explicit interface rather than a
+	 * {@code Function<InputStream, ObjectInputStream>} because the constructor throws {@link IOException}.
+	 */
+	@FunctionalInterface
+	public interface ObjectInputStreamFactory {
+		ObjectInputStream create(InputStream in) throws IOException;
+	}
+
 	/** Single-byte sentinel for null. Cannot collide with Java-serialized data (which starts 0xAC 0xED). */
 	private static final byte[] NULL_VALUE = new byte[]{0};
 
@@ -123,17 +138,34 @@ public class ValkeyCacheService implements MemcacheService {
 	/** Cold-cache sentinel options: {@code NX} plus the default TTL. Pre-built for the same reason as {@link #defaultSetOptions}. */
 	private final SetOptions defaultNxSetOptions;
 
-	/** Uses {@link #DEFAULT_EXPIRATION_SECONDS} as the fallback TTL. */
+	/** Builds the {@link ObjectInputStream} used on the read path; see {@link ObjectInputStreamFactory}. */
+	private final ObjectInputStreamFactory objectInputStreamFactory;
+
+	/** Uses {@link #DEFAULT_EXPIRATION_SECONDS} as the fallback TTL and a plain {@link ObjectInputStream}. */
 	public ValkeyCacheService(final BaseClient client) {
 		this(client, DEFAULT_EXPIRATION_SECONDS);
 	}
 
+	/** Uses {@link #DEFAULT_EXPIRATION_SECONDS} as the fallback TTL. */
+	public ValkeyCacheService(final BaseClient client, final ObjectInputStreamFactory objectInputStreamFactory) {
+		this(client, DEFAULT_EXPIRATION_SECONDS, objectInputStreamFactory);
+	}
+
+	/** Uses a plain {@link ObjectInputStream} on the read path. */
 	public ValkeyCacheService(final BaseClient client, final int defaultExpirationSeconds) {
+		this(client, defaultExpirationSeconds, ObjectInputStream::new);
+	}
+
+	public ValkeyCacheService(
+			final BaseClient client,
+			final int defaultExpirationSeconds,
+			final ObjectInputStreamFactory objectInputStreamFactory) {
 		if (defaultExpirationSeconds <= 0) {
 			throw new IllegalArgumentException("defaultExpirationSeconds must be positive, got " + defaultExpirationSeconds);
 		}
 		this.client = client;
 		this.defaultExpirationSeconds = defaultExpirationSeconds;
+		this.objectInputStreamFactory = Objects.requireNonNull(objectInputStreamFactory, "objectInputStreamFactory");
 		this.defaultSetOptions = SetOptions.builder()
 				.expiry(SetOptions.Expiry.Seconds((long) defaultExpirationSeconds))
 				.build();
@@ -157,7 +189,7 @@ private static byte[] toCacheBytes(final Object thing) {
 		return deflate(serialized);
 	}
 
-	private static Object fromCacheBytes(final byte[] bytes) {
+	private Object fromCacheBytes(final byte[] bytes) {
 		if (bytes == null || Arrays.equals(bytes, NULL_VALUE)) {
 			return null;
 		}
@@ -181,9 +213,9 @@ private static byte[] serialize(final Object thing) {
 		}
 	}
 
-	private static Object deserialize(final byte[] serialized) {
+	private Object deserialize(final byte[] serialized) {
 		try (final ByteArrayInputStream bais = new ByteArrayInputStream(serialized);
-			 final ObjectInputStream ois = new ObjectInputStream(bais)) {
+			 final ObjectInputStream ois = objectInputStreamFactory.create(bais)) {
 			return ois.readObject();
 		} catch (final IOException | ClassNotFoundException e) {
 			throw new RuntimeException("Failed to deserialize cache value", e);

src/test/java/com/googlecode/objectify/test/valkey/ValkeyCacheServiceTests.java

@@ -17,7 +17,10 @@
 import org.testcontainers.utility.DockerImageName;
 
 import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
+import java.io.ObjectStreamClass;
 import java.io.Serializable;
 import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
@@ -358,6 +361,41 @@ void readsEmptyStoredValueAsCorruptRatherThanIndexError() throws Exception {
 		}
 	}
 
+	@Test
+	void injectedFactoryRemapsClassNamesOnRead() {
+		// Simulates the shaded/un-relocated split that produces "ClassNotFoundException:
+		// streak.com.google.common..." in prod: a value serialized under one class name is read back by a
+		// process that only has the class under a different name. A ValkeyCacheService built with a custom
+		// ObjectInputStreamFactory that swaps the class descriptor must decode it — which is exactly how
+		// Streak's KeySubstitutingObjectInputStream reads cache entries written by a differently-relocated
+		// peer. The swap is done in readClassDescriptor(), not resolveClass(): since JDK 17 the latter
+		// rejects a class whose name differs from the stream descriptor with InvalidClassException.
+		final String writtenName = Alpha.class.getName();
+		final MemcacheService remapping = new ValkeyCacheService(client, in -> new ObjectInputStream(in) {
+			@Override
+			protected ObjectStreamClass readClassDescriptor() throws IOException, ClassNotFoundException {
+				final ObjectStreamClass desc = super.readClassDescriptor();
+				return desc.getName().equals(writtenName) ? ObjectStreamClass.lookup(Beta.class) : desc;
+			}
+		});
+
+		remapping.put("k", new Alpha("payload"));
+		final Object read = remapping.get("k");
+
+		assertThat(read).isInstanceOf(Beta.class);
+		assertThat(((Beta) read).s).isEqualTo("payload");
+	}
+
+	@Test
+	void factoryConstructorRoundTripsAndAppliesDefaultTtl() throws Exception {
+		// The (client, ttl, factory) constructor honors both the TTL and a plain-delegating factory.
+		final MemcacheService withFactory =
+				new ValkeyCacheService(client, 100, ObjectInputStream::new);
+		withFactory.put("k", new Alpha("v"));
+		assertThat(((Alpha) withFactory.get("k")).s).isEqualTo("v");
+		assertThat(ttlOf("k")).isGreaterThan(0L);
+	}
+
 	private static byte[] rawBytes(final String key) throws Exception {
 		final GlideString value = client.get(GlideString.gs(key.getBytes(StandardCharsets.UTF_8))).get();
 		return value == null ? null : value.getBytes();
@@ -381,6 +419,29 @@ private static String repeat(final char c, final int n) {
 		return new String(chars);
 	}
 
+	/**
+	 * {@link Alpha} and {@link Beta} share an identical serial layout and {@code serialVersionUID}, so a
+	 * stream written as {@code Alpha} deserializes into {@code Beta} once {@code resolveClass} maps the
+	 * name — the same compatibility the shaded/un-relocated Guava class pairs rely on.
+	 */
+	private static class Alpha implements Serializable {
+		private static final long serialVersionUID = 99L;
+		private final String s;
+
+		Alpha(final String s) {
+			this.s = s;
+		}
+	}
+
+	private static class Beta implements Serializable {
+		private static final long serialVersionUID = 99L;
+		private final String s;
+
+		Beta(final String s) {
+			this.s = s;
+		}
+	}
+
 	private static class Sample implements Serializable {
 		private static final long serialVersionUID = 1L;
 		private final String s;