Error Prone / NullAway support for JSpecify

build.gradle

@@ -11,11 +11,12 @@ plugins {
     id 'io.github.gradle-nexus.publish-plugin' version '1.0.0'
     id 'com.github.ben-manes.versions' version '0.51.0'
     id "me.champeau.jmh" version "0.7.3"
+    id "net.ltgt.errorprone" version '4.2.0'
 }
 
 java {
     toolchain {
-        languageVersion = JavaLanguageVersion.of(11)
+        languageVersion = JavaLanguageVersion.of(17)
     }
 }
 
@@ -75,8 +76,34 @@ dependencies {
     // this is needed for the idea jmh plugin to work correctly
     jmh 'org.openjdk.jmh:jmh-core:1.37'
     jmh 'org.openjdk.jmh:jmh-generator-annprocess:1.37'
+
+    errorprone 'com.uber.nullaway:nullaway:0.12.6'
+    errorprone 'com.google.errorprone:error_prone_core:2.37.0'
+}
+
+import net.ltgt.gradle.errorprone.CheckSeverity
+
+tasks.withType(JavaCompile) {
+    options.release = 11
+    options.errorprone {
+        disableAllChecks = true
+        check("NullAway", CheckSeverity.ERROR)
+        //
+        // end state has us with this config turned on - eg all classes
+        //
+        //option("NullAway:AnnotatedPackages", "org.dataloader")
+        option("NullAway:OnlyNullMarked", "true")
+        option("NullAway:JSpecifyMode", "true")
+    }
+    // Include to disable NullAway on test code
+    if (name.toLowerCase().contains("test")) {
+        options.errorprone {
+            disable("NullAway")
+        }
+    }
 }
 
+
 task sourcesJar(type: Jar) {
     dependsOn classes
     archiveClassifier.set('sources')

src/main/java/org/dataloader/BatchLoaderEnvironment.java

@@ -19,11 +19,11 @@
 @NullMarked
 public class BatchLoaderEnvironment {
 
-    private final Object context;
+    private final @Nullable Object context;
     private final Map<Object, Object> keyContexts;
     private final List<Object> keyContextsList;
 
-    private BatchLoaderEnvironment(Object context, List<Object> keyContextsList, Map<Object, Object> keyContexts) {
+    private BatchLoaderEnvironment(@Nullable Object context, List<Object> keyContextsList, Map<Object, Object> keyContexts) {
         this.context = context;
         this.keyContexts = keyContexts;
         this.keyContextsList = keyContextsList;
@@ -33,7 +33,6 @@ private BatchLoaderEnvironment(Object context, List<Object> keyContextsList, Map
      * Returns the overall context object provided by {@link org.dataloader.BatchLoaderContextProvider}
      *
      * @param <T> the type you would like the object to be
-     *
      * @return a context object or null if there isn't one
      */
     @SuppressWarnings("unchecked")
@@ -68,7 +67,7 @@ public static Builder newBatchLoaderEnvironment() {
     }
 
     public static class Builder {
-        private Object context;
+        private @Nullable Object context;
         private Map<Object, Object> keyContexts = Collections.emptyMap();
         private List<Object> keyContextsList = Collections.emptyList();
 

src/main/java/org/dataloader/DataLoaderRegistry.java

@@ -1,6 +1,7 @@
 package org.dataloader;
 
 import org.dataloader.annotations.PublicApi;
+import org.dataloader.impl.Assertions;
 import org.dataloader.instrumentation.ChainedDataLoaderInstrumentation;
 import org.dataloader.instrumentation.DataLoaderInstrumentation;
 import org.dataloader.instrumentation.DataLoaderInstrumentationHelper;
@@ -14,6 +15,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.ConcurrentHashMap;
 import java.util.function.Function;
@@ -142,7 +144,7 @@ private static DataLoaderOptions setInInstrumentation(DataLoaderOptions options,
      */
     public DataLoaderRegistry register(DataLoader<?, ?> dataLoader) {
         String name = dataLoader.getName();
-        assertState(name != null, () -> "The DataLoader must have a non null name");
+        Objects.requireNonNull(name, "The DataLoader must have a non null name");
         dataLoaders.put(name, nameAndInstrumentDL(name, instrumentation, dataLoader));
         return this;
     }
@@ -176,7 +178,7 @@ public DataLoaderRegistry register(String key, DataLoader<?, ?> dataLoader) {
      */
     public <K, V> DataLoader<K, V> registerAndGet(String key, DataLoader<?, ?> dataLoader) {
         dataLoaders.put(key, nameAndInstrumentDL(key, instrumentation, dataLoader));
-        return getDataLoader(key);
+        return Objects.requireNonNull(getDataLoader(key));
     }
 
     /**
@@ -251,10 +253,10 @@ public DataLoaderRegistry unregister(String key) {
      * @param key the key of the data loader
      * @param <K> the type of keys
      * @param <V> the type of values
-     * @return a data loader or null if its not present
+     * @return a data loader or null if it's not present
      */
     @SuppressWarnings("unchecked")
-    public <K, V> DataLoader<K, V> getDataLoader(String key) {
+    public <K, V> @Nullable DataLoader<K, V> getDataLoader(String key) {
         return (DataLoader<K, V>) dataLoaders.get(key);
     }
 

src/test/java/org/dataloader/DataLoaderRegistryTest.java

@@ -63,7 +63,7 @@ public void registration_works() {
         try {
             registry.register(dlUnnamed);
             Assertions.fail("Should have thrown an exception");
-        } catch (DataLoaderAssertionException ignored) {
+        } catch (NullPointerException ignored) {
         }
     }