Cleanup the otelHelper after we've finished

OctopusDeploy · Aug 16, 2022 · d7847db · d7847db
1 parent 6eea142
commit d7847db
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 5 deletions.
diff --git a/server/src/main/java/com/octopus/teamcity/opentelemetry/server/TeamCityBuildListener.java b/server/src/main/java/com/octopus/teamcity/opentelemetry/server/TeamCityBuildListener.java
@@ -164,7 +164,7 @@ private void buildFinishedOrInterrupted (SRunningBuild build) {
                     span.end();
                     var buildId = getBuildId(build);
                     otelHelper.removeSpan(buildId);
-                    if (buildId.equals(String.valueOf(getRootBuildInChain(build))))
+                    if (buildId.equals(String.valueOf(getRootBuildInChain(build).getId())))
                         otelHelperFactory.release(build.getBuildId());
                 }
             } else {

diff --git a/...va/com/octopus/teamcity/opentelemetry/server/helpers/HelperPerBuildOTELHelperFactory.java b/...va/com/octopus/teamcity/opentelemetry/server/helpers/HelperPerBuildOTELHelperFactory.java
@@ -107,6 +107,10 @@ private SpanProcessor buildGrpcSpanProcessor(Map<String, String> headers, String
 
     @Override
     public void release(Long buildId) {
-        otelHelpers.remove(buildId);
+        var helper = otelHelpers.get(buildId);
+        if (helper != null) {
+            helper.release();
+            otelHelpers.remove(buildId);
+        }
     }
 }
diff --git a/...r/src/main/java/com/octopus/teamcity/opentelemetry/server/helpers/NullOTELHelperImpl.java b/...r/src/main/java/com/octopus/teamcity/opentelemetry/server/helpers/NullOTELHelperImpl.java
@@ -35,4 +35,8 @@ public Span getSpan(String buildId) {
     @Override
     public void addAttributeToSpan(Span span, String attributeName, Object attributeValue) {
     }
+
+    @Override
+    public void release() {
+    }
 }
diff --git a/server/src/main/java/com/octopus/teamcity/opentelemetry/server/helpers/OTELHelper.java b/server/src/main/java/com/octopus/teamcity/opentelemetry/server/helpers/OTELHelper.java
@@ -37,4 +37,6 @@ public interface OTELHelper {
     Span getSpan(String buildId);
 
     void addAttributeToSpan(Span span, String attributeName, Object attributeValue);
+
+    void release();
 }
diff --git a/server/src/main/java/com/octopus/teamcity/opentelemetry/server/helpers/OTELHelperImpl.java b/server/src/main/java/com/octopus/teamcity/opentelemetry/server/helpers/OTELHelperImpl.java
@@ -13,7 +13,6 @@
 import io.opentelemetry.sdk.trace.SdkTracerProvider;
 import io.opentelemetry.sdk.trace.SpanProcessor;
 import io.opentelemetry.semconv.resource.attributes.ResourceAttributes;
-import jetbrains.buildServer.log.Loggers;
 import org.apache.log4j.Logger;
 
 import java.util.concurrent.ConcurrentHashMap;
@@ -24,19 +23,19 @@ public class OTELHelperImpl implements OTELHelper {
     private final OpenTelemetry openTelemetry;
     private final Tracer tracer;
     private final ConcurrentHashMap<String, Span> spanMap;
+    private final SdkTracerProvider sdkTracerProvider;
 
     public OTELHelperImpl(SpanProcessor spanProcessor) {
         Resource serviceNameResource = Resource
                 .create(Attributes.of(ResourceAttributes.SERVICE_NAME, PluginConstants.SERVICE_NAME));
-        SdkTracerProvider sdkTracerProvider = SdkTracerProvider.builder()
+        this.sdkTracerProvider = SdkTracerProvider.builder()
                 .setResource(Resource.getDefault().merge(serviceNameResource))
                 .addSpanProcessor(spanProcessor)
                 .build();
         this.openTelemetry = OpenTelemetrySdk.builder()
                 .setTracerProvider(sdkTracerProvider)
                 .setPropagators(ContextPropagators.create(W3CTraceContextPropagator.getInstance()))
                 .build();
-        LOG.info("OpenTelemetry plugin started.");
         this.tracer = this.openTelemetry.getTracer(PluginConstants.TRACER_INSTRUMENTATION_NAME);
         this.spanMap = new ConcurrentHashMap<>();
     }
@@ -77,4 +76,10 @@ public void addAttributeToSpan(Span span, String attributeName, Object attribute
         LOG.debug("Adding attribute to span " + attributeName + "=" + attributeValue);
         span.setAttribute(attributeName, attributeValue.toString());
     }
+
+    @Override
+    public void release() {
+        this.sdkTracerProvider.forceFlush();
+        this.sdkTracerProvider.close();
+    }
 }