Allows @AutoCodec to handle injected deserialization dependencies.

* Adds an interface, InjectingObjectCodec, taking an injected parameter in its deserialize method.
* Adds the annotation, @AutoCodec.Dependency, that can be used to indicate that a constructor parameter is a dependency instead of a normal parameter.

PiperOrigin-RevId: 180797816

7 files changed, 347 insertions, 76 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/InjectingObjectCodec.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/InjectingObjectCodec.java
new file mode 100644
index 0000000000..33da5d5992
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/InjectingObjectCodec.java
@@ -0,0 +1,48 @@
+// Copyright 2017 The Bazel Authors. All rights reserved.
+//
+// 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
+//
+//    http://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 com.google.devtools.build.lib.skyframe.serialization;
+
+import com.google.protobuf.CodedInputStream;
+import com.google.protobuf.CodedOutputStream;
+import java.io.IOException;
+
+/**
+ * Like ObjectCodec, but handles cases where deserialization requires injection of non-serializable
+ * dependencies.
+ */
+public interface InjectingObjectCodec<T, D> extends BaseCodec<T> {
+
+  /**
+   * Serializes {@code obj}, inverse of {@link #deserialize(CodedInputStream)}.
+   *
+   * @param obj the object to serialize
+   * @param codedOut the {@link CodedOutputStream} to write this object into. Implementations need
+   *     not call {@link CodedOutputStream#flush()}, this should be handled by the caller.
+   * @throws SerializationException on failure to serialize
+   * @throws IOException on {@link IOException} during serialization
+   */
+  void serialize(T obj, CodedOutputStream codedOut) throws SerializationException, IOException;
+
+  /**
+   * Deserializes from {@code codedIn}, inverse of {@link #serialize(Object, CodedOutputStream)}.
+   *
+   * @param dependency the non-serializable, injected dependency
+   * @param codedIn the {@link CodedInputStream} to read the serialized object from
+   * @return the object deserialized from {@code codedIn}
+   * @throws SerializationException on failure to deserialize
+   * @throws IOException on {@link IOException} during deserialization
+   */
+  T deserialize(D dependency, CodedInputStream codedIn) throws SerializationException, IOException;
+}
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/InjectingObjectCodecAdapter.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/InjectingObjectCodecAdapter.java
new file mode 100644
index 0000000000..a14257395f
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/InjectingObjectCodecAdapter.java
@@ -0,0 +1,53 @@
+// Copyright 2017 The Bazel Authors. All rights reserved.
+//
+// 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
+//
+//    http://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 com.google.devtools.build.lib.skyframe.serialization;
+
+import com.google.protobuf.CodedInputStream;
+import com.google.protobuf.CodedOutputStream;
+import java.io.IOException;
+
+/** Adapts an InjectingObjectCodec to an ObjectCodec. */
+public class InjectingObjectCodecAdapter<T, D> implements ObjectCodec<T> {
+
+  private final InjectingObjectCodec<T, D> codec;
+  private final D dependency;
+
+  /**
+   * Creates an ObjectCodec from an InjectingObjectCodec.
+   *
+   * @param codec underlying codec to delegate to
+   * @param dependency object forwarded to {@code codec.deserialize}
+   */
+  public InjectingObjectCodecAdapter(InjectingObjectCodec<T, D> codec, D dependency) {
+    this.codec = codec;
+    this.dependency = dependency;
+  }
+
+  @Override
+  public Class<T> getEncodedClass() {
+    return codec.getEncodedClass();
+  }
+
+  @Override
+  public void serialize(T obj, CodedOutputStream codedOut)
+      throws SerializationException, IOException {
+    codec.serialize(obj, codedOut);
+  }
+
+  @Override
+  public T deserialize(CodedInputStream codedIn) throws SerializationException, IOException {
+    return codec.deserialize(dependency, codedIn);
+  }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/AutoCodec.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/AutoCodec.java
index 2d67bedf6c..638b18897a 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/AutoCodec.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/AutoCodec.java
@@ -78,5 +78,25 @@ public @interface AutoCodec {
   @Target(ElementType.CONSTRUCTOR)
   public static @interface Constructor {}
 
+  /**
+   * Marks a specific constructor parameter as a dependency when using the {@code CONSTRUCTOR}
+   * strategy.
+   *
+   * <p>When a constructor selected for the {@code CONSTRUCTOR} strategy has one of its parameters
+   * tagged {@code @Dependency}, {@code @AutoCodec} generates an {@link
+   * com.google.devtools.build.lib.skyframe.serialization.InjectingObjectCodec} instead of the usual
+   * {@link com.google.devtools.build.lib.skyframe.serialization.ObjectCodec} with the dependency
+   * type parameter matching the tagged parameter type.
+   *
+   * <p>At deserialization, the {@code @Dependency} tagged parameter will be forwarded from the
+   * {@code dependency} parameter of {@link
+   * com.google.devtools.build.lib.skyframe.serialization.InjectingObjectCodec#deserialize}.
+   *
+   * <p>A compiler error will result if more than one constructor parameter has the
+   * {@code @Dependency} annotation.
+   */
+  @Target(ElementType.PARAMETER)
+  public static @interface Dependency {}
+
   Strategy strategy() default Strategy.CONSTRUCTOR;
 }
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/AutoCodecProcessor.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/AutoCodecProcessor.java
index 7c7a1d1fb2..0c4d34bef3 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/AutoCodecProcessor.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/AutoCodecProcessor.java
@@ -33,8 +33,10 @@ import java.lang.reflect.Field;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
+import javax.annotation.Nullable;
 import javax.annotation.processing.AbstractProcessor;
 import javax.annotation.processing.ProcessingEnvironment;
 import javax.annotation.processing.Processor;
@@ -94,24 +96,24 @@ public class AutoCodecProcessor extends AbstractProcessor {
     for (Element element : roundEnv.getElementsAnnotatedWith(AutoCodecUtil.ANNOTATION)) {
       AutoCodec annotation = element.getAnnotation(AutoCodecUtil.ANNOTATION);
       TypeElement encodedType = (TypeElement) element;
-      TypeSpec.Builder codecClassBuilder = AutoCodecUtil.initializeCodecClassBuilder(encodedType);
-      codecClassBuilder.addMethod(
-          AutoCodecUtil.initializeGetEncodedClassMethod(encodedType)
-              .addStatement("return $T.class", TypeName.get(encodedType.asType()))
-              .build());
+      TypeSpec.Builder codecClassBuilder = null;
       switch (annotation.strategy()) {
         case CONSTRUCTOR:
-          buildClassWithConstructorStrategy(codecClassBuilder, encodedType);
+          codecClassBuilder = buildClassWithConstructorStrategy(encodedType);
           break;
         case PUBLIC_FIELDS:
-          buildClassWithPublicFieldsStrategy(codecClassBuilder, encodedType);
+          codecClassBuilder = buildClassWithPublicFieldsStrategy(encodedType);
           break;
         case POLYMORPHIC:
-          buildClassWithPolymorphicStrategy(codecClassBuilder, encodedType);
+          codecClassBuilder = buildClassWithPolymorphicStrategy(encodedType);
           break;
         default:
           throw new IllegalArgumentException("Unknown strategy: " + annotation.strategy());
       }
+      codecClassBuilder.addMethod(
+          AutoCodecUtil.initializeGetEncodedClassMethod(encodedType)
+              .addStatement("return $T.class", TypeName.get(encodedType.asType()))
+              .build());
       String packageName =
           env.getElementUtils().getPackageOf(encodedType).getQualifiedName().toString();
       try {
@@ -129,8 +131,63 @@ public class AutoCodecProcessor extends AbstractProcessor {
     return true;
   }
 
-  private void buildClassWithConstructorStrategy(
-      TypeSpec.Builder codecClassBuilder, TypeElement encodedType) {
+  private TypeSpec.Builder buildClassWithConstructorStrategy(TypeElement encodedType) {
+    ExecutableElement constructor = selectConstructorForConstructorStrategy(encodedType);
+    PartitionedParameters parameters = isolateDependency(constructor);
+
+    TypeSpec.Builder codecClassBuilder =
+        AutoCodecUtil.initializeCodecClassBuilder(encodedType, parameters.dependency);
+
+    initializeUnsafeOffsets(codecClassBuilder, encodedType, parameters.fields);
+
+    codecClassBuilder.addMethod(
+        buildSerializeMethodWithConstructor(encodedType, parameters.fields));
+
+    MethodSpec.Builder deserializeBuilder =
+        AutoCodecUtil.initializeDeserializeMethodBuilder(encodedType, parameters.dependency);
+    buildDeserializeBody(deserializeBuilder, parameters.fields);
+    addReturnNew(deserializeBuilder, encodedType, constructor);
+    codecClassBuilder.addMethod(deserializeBuilder.build());
+
+    return codecClassBuilder;
+  }
+
+  private static class PartitionedParameters {
+    /** Non-dependency parameters. */
+    List<VariableElement> fields;
+    /**
+     * Parameter having the {@link AutoCodec.Dependency} annotation.
+     *
+     * <p>Null if no such parameter exists.
+     */
+    @Nullable VariableElement dependency;
+  }
+
+  /** Separates any dependency from the constructor parameters. */
+  private static PartitionedParameters isolateDependency(ExecutableElement constructor) {
+    Map<Boolean, List<VariableElement>> splitParameters =
+        constructor
+            .getParameters()
+            .stream()
+            .collect(
+                Collectors.partitioningBy(
+                    p -> p.getAnnotation(AutoCodec.Dependency.class) != null));
+    PartitionedParameters result = new PartitionedParameters();
+    result.fields = splitParameters.get(Boolean.FALSE);
+    List<VariableElement> dependencies = splitParameters.get(Boolean.TRUE);
+    if (dependencies.size() > 1) {
+      throw new IllegalArgumentException(
+          ((TypeElement) constructor.getEnclosingElement()).getQualifiedName()
+              + " constructor has multiple Dependency annotations.");
+    }
+    if (!dependencies.isEmpty()) {
+      result.dependency = dependencies.get(0);
+    }
+    return result;
+  }
+
+  private static ExecutableElement selectConstructorForConstructorStrategy(
+      TypeElement encodedType) {
     List<ExecutableElement> constructors =
         ElementFilter.constructorsIn(encodedType.getEnclosedElements());
     ImmutableList<ExecutableElement> markedConstructors =
@@ -138,7 +195,6 @@ public class AutoCodecProcessor extends AbstractProcessor {
             .stream()
             .filter(c -> c.getAnnotation(AutoCodec.Constructor.class) != null)
             .collect(toImmutableList());
-    ExecutableElement constructor = null;
     if (markedConstructors.isEmpty()) {
       // If nothing is marked, see if there is a unique constructor.
       if (constructors.size() > 1) {
@@ -147,22 +203,13 @@ public class AutoCodecProcessor extends AbstractProcessor {
                 + " has multiple constructors but no Constructor annotation.");
       }
       // In Java, every class has at least one constructor, so this never fails.
-      constructor = constructors.get(0);
-    } else if (markedConstructors.size() == 1) {
-      constructor = markedConstructors.get(0);
-    } else {
-      throw new IllegalArgumentException(
-          encodedType.getQualifiedName() + " has multiple Constructor annotations.");
+      return constructors.get(0);
     }
-    List<? extends VariableElement> constructorParameters = constructor.getParameters();
-    initializeUnsafeOffsets(codecClassBuilder, encodedType, constructorParameters);
-    codecClassBuilder.addMethod(
-        buildSerializeMethodWithConstructor(encodedType, constructorParameters));
-    MethodSpec.Builder deserializeBuilder =
-        AutoCodecUtil.initializeDeserializeMethodBuilder(encodedType);
-    buildDeserializeBody(deserializeBuilder, constructorParameters);
-    addReturnNew(deserializeBuilder, encodedType, constructor, constructorParameters);
-    codecClassBuilder.addMethod(deserializeBuilder.build());
+    if (markedConstructors.size() == 1) {
+      return markedConstructors.get(0);
+    }
+    throw new IllegalArgumentException(
+        encodedType.getQualifiedName() + " has multiple Constructor annotations.");
   }
 
   private MethodSpec buildSerializeMethodWithConstructor(
@@ -206,19 +253,20 @@ public class AutoCodecProcessor extends AbstractProcessor {
     return serializeBuilder.build();
   }
 
-  private void buildClassWithPublicFieldsStrategy(
-      TypeSpec.Builder codecClassBuilder, TypeElement encodedType) {
-    List<? extends VariableElement> publicFields =
+  private TypeSpec.Builder buildClassWithPublicFieldsStrategy(TypeElement encodedType) {
+    TypeSpec.Builder codecClassBuilder = AutoCodecUtil.initializeCodecClassBuilder(encodedType);
+    ImmutableList<? extends VariableElement> publicFields =
         ElementFilter.fieldsIn(env.getElementUtils().getAllMembers(encodedType))
             .stream()
             .filter(this::isPublicField)
-            .collect(Collectors.toList());
+            .collect(toImmutableList());
     codecClassBuilder.addMethod(buildSerializeMethodWithPublicFields(encodedType, publicFields));
     MethodSpec.Builder deserializeBuilder =
         AutoCodecUtil.initializeDeserializeMethodBuilder(encodedType);
     buildDeserializeBody(deserializeBuilder, publicFields);
     addInstantiatePopulateFieldsAndReturn(deserializeBuilder, encodedType, publicFields);
     codecClassBuilder.addMethod(deserializeBuilder.build());
+    return codecClassBuilder;
   }
 
   private boolean isPublicField(VariableElement element) {
@@ -290,18 +338,18 @@ public class AutoCodecProcessor extends AbstractProcessor {
    * <p>Used by the {@link AutoCodec.Strategy.CONSTRUCTOR} strategy.
    */
   private static void addReturnNew(
-      MethodSpec.Builder builder,
-      TypeElement type,
-      ExecutableElement constructor,
-      List<? extends VariableElement> parameters) {
+      MethodSpec.Builder builder, TypeElement type, ExecutableElement constructor) {
     List<? extends TypeMirror> allThrown = constructor.getThrownTypes();
     if (!allThrown.isEmpty()) {
       builder.beginControlFlow("try");
     }
-    builder.addStatement(
-        "return new $T($L)",
-        TypeName.get(type.asType()),
-        parameters.stream().map(p -> p.getSimpleName() + "_").collect(Collectors.joining(", ")));
+    String parameters =
+        constructor
+            .getParameters()
+            .stream()
+            .map(AutoCodecProcessor::handleFromParameter)
+            .collect(Collectors.joining(", "));
+    builder.addStatement("return new $T($L)", TypeName.get(type.asType()), parameters);
     if (!allThrown.isEmpty()) {
       for (TypeMirror thrown : allThrown) {
         builder.nextControlFlow("catch ($T e)", TypeName.get(thrown));
@@ -315,6 +363,18 @@ public class AutoCodecProcessor extends AbstractProcessor {
   }
 
   /**
+   * Coverts a constructor parameter to a String representing its handle within deserialize.
+   *
+   * <p>Uses the handle {@code dependency} for any parameter with the {@link AutoCodec.Dependency}
+   * annotation.
+   */
+  private static String handleFromParameter(VariableElement parameter) {
+    return parameter.getAnnotation(AutoCodec.Dependency.class) != null
+        ? "dependency"
+        : (parameter.getSimpleName() + "_");
+  }
+
+  /**
    * Invokes the constructor, populates public fields and returns the value.
    *
    * <p>Used by the {@link AutoCodec.Strategy.PUBLIC_FIELDS} strategy.
@@ -375,7 +435,7 @@ public class AutoCodecProcessor extends AbstractProcessor {
    *
    * <p>Throws IllegalArgumentException if no such field is found.
    */
-  private VariableElement getFieldByName(TypeElement type, String name) {
+  private static VariableElement getFieldByName(TypeElement type, String name) {
     return ElementFilter.fieldsIn(type.getEnclosedElements())
         .stream()
         .filter(f -> f.getSimpleName().contentEquals(name))
@@ -386,14 +446,15 @@ public class AutoCodecProcessor extends AbstractProcessor {
                     type.getQualifiedName() + ": no field with name matching " + name));
   }
 
-  private static void buildClassWithPolymorphicStrategy(
-      TypeSpec.Builder codecClassBuilder, TypeElement encodedType) {
+  private static TypeSpec.Builder buildClassWithPolymorphicStrategy(TypeElement encodedType) {
     if (!encodedType.getModifiers().contains(Modifier.ABSTRACT)) {
       throw new IllegalArgumentException(
           encodedType + " is not abstract, but POLYMORPHIC was selected as the strategy.");
     }
+    TypeSpec.Builder codecClassBuilder = AutoCodecUtil.initializeCodecClassBuilder(encodedType);
     codecClassBuilder.addMethod(buildPolymorphicSerializeMethod(encodedType));
     codecClassBuilder.addMethod(buildPolymorphicDeserializeMethod(encodedType));
+    return codecClassBuilder;
   }
 
   private static MethodSpec buildPolymorphicSerializeMethod(TypeElement encodedType) {
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/AutoCodecUtil.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/AutoCodecUtil.java
index 7dbc42952f..e958c3b60b 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/AutoCodecUtil.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/AutoCodecUtil.java
@@ -15,6 +15,7 @@
 package com.google.devtools.build.lib.skyframe.serialization.autocodec;
 
 import com.google.common.collect.ImmutableList;
+import com.google.devtools.build.lib.skyframe.serialization.InjectingObjectCodec;
 import com.google.devtools.build.lib.skyframe.serialization.ObjectCodec;
 import com.google.devtools.build.lib.skyframe.serialization.SerializationException;
 import com.google.protobuf.CodedInputStream;
@@ -26,9 +27,11 @@ import com.squareup.javapoet.TypeName;
 import com.squareup.javapoet.TypeSpec;
 import java.io.IOException;
 import java.util.stream.Collectors;
+import javax.annotation.Nullable;
 import javax.lang.model.element.Element;
 import javax.lang.model.element.Modifier;
 import javax.lang.model.element.TypeElement;
+import javax.lang.model.element.VariableElement;
 
 /** Static utilities for AutoCodec processors. */
 class AutoCodecUtil {
@@ -36,11 +39,37 @@ class AutoCodecUtil {
   public static final String GENERATED_CLASS_NAME_SUFFIX = "AutoCodec";
   static final Class<AutoCodec> ANNOTATION = AutoCodec.class;
 
+  /**
+   * Initializes a builder for a class implementing {@link ObjectCodec}.
+   *
+   * @param encodedType type being serialized
+   */
   static TypeSpec.Builder initializeCodecClassBuilder(TypeElement encodedType) {
-    return TypeSpec.classBuilder(getCodecName(encodedType))
-        .addSuperinterface(
-            ParameterizedTypeName.get(
-                ClassName.get(ObjectCodec.class), TypeName.get(encodedType.asType())));
+    return initializeCodecClassBuilder(encodedType, null);
+  }
+
+  /**
+   * Initializes a builder for a class of the appropriate type.
+   *
+   * <p>If the dependency is non-null, then the type is {@link InjectingObjectCodec} otherwise
+   * {@link ObjectCodec}.
+   *
+   * @param encodedType type being serialized
+   * @param dependency type being injected or null
+   */
+  static TypeSpec.Builder initializeCodecClassBuilder(
+      TypeElement encodedType, @Nullable VariableElement dependency) {
+    TypeSpec.Builder builder = TypeSpec.classBuilder(getCodecName(encodedType));
+    if (dependency == null) {
+      return builder.addSuperinterface(
+          ParameterizedTypeName.get(
+              ClassName.get(ObjectCodec.class), TypeName.get(encodedType.asType())));
+    }
+    return builder.addSuperinterface(
+        ParameterizedTypeName.get(
+            ClassName.get(InjectingObjectCodec.class),
+            TypeName.get(encodedType.asType()),
+            TypeName.get(dependency.asType())));
   }
 
   static MethodSpec.Builder initializeGetEncodedClassMethod(TypeElement encodedType) {
@@ -63,14 +92,33 @@ class AutoCodecUtil {
         .addException(IOException.class);
   }
 
+  /** Initializes {@link ObjectCodec#deserialize}. */
   static MethodSpec.Builder initializeDeserializeMethodBuilder(TypeElement encodedType) {
-    return MethodSpec.methodBuilder("deserialize")
-        .addModifiers(Modifier.PUBLIC)
-        .returns(TypeName.get(encodedType.asType()))
-        .addParameter(CodedInputStream.class, "codedIn")
-        .addAnnotation(Override.class)
-        .addException(SerializationException.class)
-        .addException(IOException.class);
+    return initializeDeserializeMethodBuilder(encodedType, null);
+  }
+
+  /**
+   * Initializes the appropriate deserialize method based on presence of dependency.
+   *
+   * <p>{@link InjectingObjectCodec#deserialize} if dependency is non-null and {@link
+   * ObjectCodec#deserialize} otherwise.
+   *
+   * @param encodedType type being serialized
+   * @param dependency type being injected
+   */
+  static MethodSpec.Builder initializeDeserializeMethodBuilder(
+      TypeElement encodedType, @Nullable VariableElement dependency) {
+    MethodSpec.Builder builder =
+        MethodSpec.methodBuilder("deserialize")
+            .addModifiers(Modifier.PUBLIC)
+            .returns(TypeName.get(encodedType.asType()))
+            .addAnnotation(Override.class)
+            .addException(SerializationException.class)
+            .addException(IOException.class);
+    if (dependency != null) {
+      builder.addParameter(TypeName.get(dependency.asType()), "dependency");
+    }
+    return builder.addParameter(CodedInputStream.class, "codedIn");
   }
 
   /**
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/BUILD b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/BUILD
index b6bc9b9fc3..3b0d365196 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/BUILD
@@ -55,6 +55,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization",
         "//third_party:auto_service",
         "//third_party:guava",
+        "//third_party:jsr305",
         "//third_party/java/javapoet",
         "//third_party/protobuf:protobuf_java",
     ],
@@ -68,6 +69,7 @@ java_library(
         ":autocodec-annotation",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization",
         "//third_party:guava",
+        "//third_party:jsr305",
         "//third_party/java/javapoet",
         "//third_party/protobuf:protobuf_java",
     ],
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/Marshallers.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/Marshallers.java
index e3c7898c58..7a01498151 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/Marshallers.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/autocodec/Marshallers.java
@@ -22,6 +22,8 @@ import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.collect.ImmutableSortedSet;
 import com.google.common.collect.Maps;
+import com.google.devtools.build.lib.skyframe.serialization.InjectingObjectCodec;
+import com.google.devtools.build.lib.skyframe.serialization.ObjectCodec;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.Marshaller.Context;
 import com.google.devtools.build.lib.skyframe.serialization.strings.StringCodecs;
 import com.google.protobuf.GeneratedMessage;
@@ -32,7 +34,10 @@ import java.util.List;
 import java.util.Map;
 import java.util.regex.Pattern;
 import javax.annotation.processing.ProcessingEnvironment;
+import javax.lang.model.element.Element;
 import javax.lang.model.element.ElementKind;
+import javax.lang.model.element.Modifier;
+import javax.lang.model.element.TypeElement;
 import javax.lang.model.type.DeclaredType;
 import javax.lang.model.type.TypeMirror;
 
@@ -61,31 +66,16 @@ class Marshallers {
   }
 
   private Marshaller getMatchingMarshaller(DeclaredType type) {
-    return marshallers.stream().filter(marshaller -> marshaller.matches(type)).findFirst().get();
+    return marshallers
+        .stream()
+        .filter(marshaller -> marshaller.matches(type))
+        .findFirst()
+        .orElseThrow(
+            () ->
+                new IllegalArgumentException(
+                    "No marshaller for: " + ((TypeElement) type.asElement()).getQualifiedName()));
   }
 
-  private static final Marshaller CODEC_MARSHALLER =
-      new Marshaller() {
-        @Override
-        public boolean matches(DeclaredType type) {
-          // TODO(shahan): check for getCodec or CODEC.
-          // CODEC is the final fallback for all Marshallers so this returns true.
-          return true;
-        }
-
-        @Override
-        public void addSerializationCode(Context context) {
-          context.builder.addStatement(
-              "$T.CODEC.serialize($L, codedOut)", context.getTypeName(), context.name);
-        }
-
-        @Override
-        public void addDeserializationCode(Context context) {
-          context.builder.addStatement(
-              "$L = $T.CODEC.deserialize(codedIn)", context.name, context.getTypeName());
-        }
-      };
-
   private final Marshaller enumMarshaller =
       new Marshaller() {
         @Override
@@ -450,6 +440,39 @@ class Marshallers {
         }
       };
 
+  private final Marshaller codecMarshaller =
+      new Marshaller() {
+        @Override
+        public boolean matches(DeclaredType type) {
+          return getCodec(type).isPresent();
+        }
+
+        @Override
+        public void addSerializationCode(Context context) {
+          context.builder.addStatement(
+              "$T.CODEC.serialize($L, codedOut)", context.getTypeName(), context.name);
+        }
+
+        @Override
+        public void addDeserializationCode(Context context) {
+          TypeMirror codecType = getCodec(context.type).get().asType();
+          if (isSubtypeErased(codecType, ObjectCodec.class)) {
+            context.builder.addStatement(
+                "$L = $T.CODEC.deserialize(codedIn)", context.name, context.getTypeName());
+          } else if (isSubtypeErased(codecType, InjectingObjectCodec.class)) {
+            context.builder.addStatement(
+                "$L = $T.CODEC.deserialize(dependency, codedIn)",
+                context.name,
+                context.getTypeName());
+          } else {
+            throw new IllegalArgumentException(
+                "CODEC field of "
+                    + ((TypeElement) context.type.asElement()).getQualifiedName()
+                    + " is neither ObjectCodec nor InjectingCodec");
+          }
+        }
+      };
+
   private final ImmutableList<Marshaller> marshallers =
       ImmutableList.of(
           enumMarshaller,
@@ -462,7 +485,7 @@ class Marshallers {
           multimapMarshaller,
           patternMarshaller,
           protoMarshaller,
-          CODEC_MARSHALLER);
+          codecMarshaller);
 
   /** True when {@code type} has the same type as {@code clazz}. */
   private boolean matchesType(TypeMirror type, Class<?> clazz) {
@@ -480,8 +503,24 @@ class Marshallers {
         .isSameType(env.getTypeUtils().erasure(type), env.getTypeUtils().erasure(getType(clazz)));
   }
 
+  /** True when erasure of {@code type} is a subtype of the erasure of {@code clazz}. */
+  private boolean isSubtypeErased(TypeMirror type, Class<?> clazz) {
+    return env.getTypeUtils()
+        .isSubtype(env.getTypeUtils().erasure(type), env.getTypeUtils().erasure(getType(clazz)));
+  }
+
   /** Returns the TypeMirror corresponding to {@code clazz}. */
   private TypeMirror getType(Class<?> clazz) {
     return env.getElementUtils().getTypeElement((clazz.getCanonicalName())).asType();
   }
+
+  private static java.util.Optional<? extends Element> getCodec(DeclaredType type) {
+    return type.asElement()
+        .getEnclosedElements()
+        .stream()
+        .filter(t -> t.getModifiers().contains(Modifier.STATIC))
+        .filter(t -> t.getSimpleName().contentEquals("CODEC"))
+        .filter(t -> t.getKind() == ElementKind.FIELD)
+        .findAny();
+  }
 }