Merge pull request #33918: Remove type suppressions from MetricsPusher and MetricsSink

Author: kennknowles

runners/core-java/src/main/java/org/apache/beam/runners/core/metrics/MetricsPusher.java

@@ -35,9 +35,6 @@
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 /** Component that regularly merges metrics and pushes them to a metrics sink. */
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 public class MetricsPusher implements Serializable {
 
   private MetricsSink metricsSink;
@@ -76,7 +73,7 @@ public void start() {
 
   private void tearDown() {
     pushMetrics();
-    if (!scheduledFuture.isCancelled()) {
+    if (scheduledFuture != null && !scheduledFuture.isCancelled()) {
       scheduledFuture.cancel(true);
     }
   }

runners/extensions-java/metrics/src/main/java/org/apache/beam/runners/extensions/metrics/MetricsGraphiteSink.java

@@ -31,6 +31,7 @@
 import org.apache.beam.sdk.metrics.MetricResult;
 import org.apache.beam.sdk.metrics.MetricsOptions;
 import org.apache.beam.sdk.metrics.MetricsSink;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
  * Sink to push metrics to Graphite. Graphite requires a timestamp. So metrics are reported with the
@@ -40,9 +41,6 @@
  * {@code beam.counter.throughput.nbRecords.attempted.value} Or {@code
  * beam.distribution.throughput.nbRecordsPerSec.attempted.mean}
  */
-@SuppressWarnings({
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 public class MetricsGraphiteSink implements MetricsSink {
   private static final Charset UTF_8 = Charset.forName("UTF-8");
   private static final Pattern WHITESPACE = Pattern.compile("[\\s]+");
@@ -103,7 +101,8 @@ public String toString() {
       try {
         messagePayload.append(createCommittedMessage());
       } catch (UnsupportedOperationException e) {
-        if (!e.getMessage().contains("committed metrics")) {
+        @Nullable String message = e.getMessage();
+        if (message != null && !message.contains("committed metrics")) {
           throw e;
         }
       }

runners/extensions-java/metrics/src/main/java/org/apache/beam/runners/extensions/metrics/MetricsHttpSink.java

@@ -41,10 +41,6 @@
 import org.apache.beam.sdk.metrics.MetricsSink;
 
 /** HTTP Sink to push metrics in a POST HTTP request. */
-@SuppressWarnings({
-  "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
-  "nullness" // TODO(https://github.com/apache/beam/issues/20497)
-})
 public class MetricsHttpSink implements MetricsSink {
   private final String urlString;
   private final ObjectMapper objectMapper = new ObjectMapper();
@@ -112,6 +108,7 @@ public MetricKeySerializer(Class<MetricKey> t) {
       super(t);
     }
 
+    @SuppressWarnings("nullness") // gen.writeObjectField expects nulls but is unannotated
     public void inline(MetricKey value, JsonGenerator gen, SerializerProvider provider)
         throws IOException {
       gen.writeObjectField("name", value.metricName());
@@ -131,32 +128,41 @@ public void serialize(MetricKey value, JsonGenerator gen, SerializerProvider pro
    * JSON serializer for {@link MetricResult}; conform to an older format where the {@link MetricKey
    * key's} {@link MetricName name} and "step" (ptransform) are inlined.
    */
-  public static class MetricResultSerializer extends StdSerializer<MetricResult> {
+  public static class MetricResultSerializer extends StdSerializer<MetricResult<?>> {
     private final MetricKeySerializer keySerializer;
 
-    public MetricResultSerializer(Class<MetricResult> t) {
+    public MetricResultSerializer(Class<MetricResult<?>> t) {
       super(t);
       keySerializer = new MetricKeySerializer(MetricKey.class);
     }
 
     @Override
-    public void serialize(MetricResult value, JsonGenerator gen, SerializerProvider provider)
+    public void serialize(MetricResult<?> value, JsonGenerator gen, SerializerProvider provider)
         throws IOException {
       gen.writeStartObject();
+      writeAttemptedAndCommitted(value, gen);
+      keySerializer.inline(value.getKey(), gen, provider);
+      gen.writeEndObject();
+    }
+
+    @SuppressWarnings("nullness") // gen.writeObjectField expects nulls but is unannotated
+    private void writeAttemptedAndCommitted(MetricResult<?> value, JsonGenerator gen)
+        throws IOException {
       gen.writeObjectField("attempted", value.getAttempted());
       if (value.hasCommitted()) {
         gen.writeObjectField("committed", value.getCommitted());
       }
-      keySerializer.inline(value.getKey(), gen, provider);
-      gen.writeEndObject();
     }
   }
 
   private String serializeMetrics(MetricQueryResults metricQueryResults) throws Exception {
     SimpleModule module = new JodaModule();
     module.addSerializer(new MetricNameSerializer(MetricName.class));
     module.addSerializer(new MetricKeySerializer(MetricKey.class));
-    module.addSerializer(new MetricResultSerializer(MetricResult.class));
+    // This odd cast converts from rawtype Class<MetricResult> to Class<MetricResult<?>>
+    // so the rest of the file can be properly typed
+    module.addSerializer(
+        new MetricResultSerializer((Class<MetricResult<?>>) (Object) MetricResult.class));
     objectMapper.registerModule(module);
     objectMapper.disable(SerializationFeature.WRITE_DATES_AS_TIMESTAMPS);
     objectMapper.configure(MapperFeature.SORT_PROPERTIES_ALPHABETICALLY, true);
@@ -171,6 +177,7 @@ private String serializeMetrics(MetricQueryResults metricQueryResults) throws Ex
       result = objectMapper.writeValueAsString(metricQueryResults);
     } catch (JsonMappingException exception) {
       if ((exception.getCause() instanceof UnsupportedOperationException)
+          && exception.getCause().getMessage() != null
           && exception.getCause().getMessage().contains("committed metrics")) {
         filterProvider.removeFilter("committedMetrics");
         filter = SimpleBeanPropertyFilter.serializeAllExcept("committed");