/**
 * Copyright 2017 Tristan Burch.
 * <p>
 * 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
 * <p>
 * http://www.apache.org/licenses/LICENSE-2.0
 * <p>
 * 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.lowtuna.dropwizard.grpc.metrics;

import io.grpc.ForwardingServerCallListener;
import io.grpc.MethodDescriptor;
import io.grpc.ServerCall;

import java.time.Duration;
import java.time.Instant;
import java.util.Objects;

public class InstrumentedServerCallListener<ReqT> extends ForwardingServerCallListener.SimpleForwardingServerCallListener<ReqT> {
    private final MethodDescriptor methodDescriptor;
    private final ServerMetrics serverMetrics;

    private Instant start;

    InstrumentedServerCallListener(ServerCall.Listener<ReqT> delegate, MethodDescriptor methodDescriptor, ServerMetrics serverMetrics) {
       super(delegate);
        this.methodDescriptor = methodDescriptor;
        this.serverMetrics = serverMetrics;
    }

    @Override
    public void onReady() {
        if (Objects.isNull(start)) {
            start = Instant.now();
            serverMetrics.incActiveRequests();
        }
        super.onReady();
    }

    @Override
    public void onMessage(ReqT request) {
        if (!methodDescriptor.getType().clientSendsOneMessage()) {
            serverMetrics.markMessageReceived(methodDescriptor);
        }
        super.onMessage(request);
    }

    @Override
    public void onComplete() {
        serverMetrics.decActiveRequests();
        super.onComplete();
        if (Objects.nonNull(start)) {
            Duration elapsed = Duration.between(start, Instant.now());
            serverMetrics.markMessageComplete(elapsed, methodDescriptor);
        }
    }

    @Override
    public void onCancel() {
        serverMetrics.decActiveRequests();
        super.onCancel();
        if (Objects.nonNull(start)) {
            Duration elapsed = Duration.between(start, Instant.now());
            serverMetrics.markMessageCanceled(elapsed, methodDescriptor);
        }
    }
}

@@ -822,23 +822,49 @@
 va.util.
-Objects
+concurrent.atomic.AtomicReference
 ;%0A%0Apubli
@@ -1105,21 +1105,70 @@
 ate 
+final AtomicReference%3C
 Instant
+%3E
  start
+ = new AtomicReference%3C%3E()
 ;%0A%0A 
@@ -1306,24 +1306,25 @@
 s) %7B%0A       
+ 
 super(delega
@@ -1483,56 +1483,33 @@
     
-if (Objects.isNull(start)) %7B%0A            start =
+start.compareAndSet(null,
  Ins
@@ -1518,22 +1518,19 @@
 nt.now()
+)
 ;%0A
-    
         
@@ -1564,26 +1564,16 @@
 ests();%0A
-        %7D%0A
         
@@ -1892,43 +1892,18 @@
     
-serverMetrics.decActiveRequests();%0A
+try %7B%0A    
     
@@ -1938,106 +1938,64 @@
     
-if (Objects.nonNull(start)) %7B%0A            Duration elapsed = Duration.between(start, Instant.now()
+%7D finally %7B%0A            serverMetrics.decActiveRequests(
 );%0A 
@@ -2039,23 +2039,60 @@
 omplete(
-elapsed
+Duration.between(start.get(), Instant.now())
 , method
@@ -2176,175 +2176,108 @@
     
-serverMetrics.decActiveRequests();%0A        super.onCancel();%0A        if (Objects.nonNull(start)) %7B%0A            Duration elapsed = Duration.between(start, Instant.now()
+try %7B%0A            super.onCancel();%0A        %7D finally %7B%0A            serverMetrics.decActiveRequests(
 );%0A 
@@ -2325,15 +2325,52 @@
 led(
-elapsed
+Duration.between(start.get(), Instant.now())
 , me

make sure that methods are called using try/finally