From c31746bddb43c951aef5e1b67a090b645bc11334 Mon Sep 17 00:00:00 2001
From: Ian Adams <ian.adams@observiq.com>
Date: Mon, 11 Nov 2024 15:11:52 -0500
Subject: [PATCH] [receiver/azureeventhub] Close storage extension client
 during component shutdown

---
 .chloggen/aeh-storageclient-shutdown.yaml     | 27 +++++++++++++++++++
 .../azureeventhubreceiver/eventhubhandler.go  | 20 ++++++++++----
 2 files changed, 42 insertions(+), 5 deletions(-)
 create mode 100644 .chloggen/aeh-storageclient-shutdown.yaml

diff --git a/.chloggen/aeh-storageclient-shutdown.yaml b/.chloggen/aeh-storageclient-shutdown.yaml
new file mode 100644
index 000000000000..d5b09a5e52ac
--- /dev/null
+++ b/.chloggen/aeh-storageclient-shutdown.yaml
@@ -0,0 +1,27 @@
+# Use this changelog template to create an entry for release notes.
+
+# One of 'breaking', 'deprecation', 'new_component', 'enhancement', 'bug_fix'
+change_type: bug_fix
+
+# The name of the component, or a single word describing the area of concern, (e.g. filelogreceiver)
+component: receiver/azureeventhub
+
+# A brief description of the change.  Surround your text with quotes ("") if it needs to start with a backtick (`).
+note: When using a storage extension, the component will call Close on the client during component shutdown. This fixes a bug that resulted in a file potentially remaining locked after component shutdown.
+
+# Mandatory: One or more tracking issues related to the change. You can use the PR number here if no issue exists.
+issues: [36238]
+
+# (Optional) One or more lines of additional information to render under the primary note.
+# These lines will be padded with 2 spaces and then inserted directly into the document.
+# Use pipe (|) for multiline entries.
+subtext:
+
+# If your change doesn't affect end users or the exported elements of any package,
+# you should instead start your pull request title with [chore] or use the "Skip Changelog" label.
+# Optional: The change log or logs in which this entry should be included.
+# e.g. '[user]' or '[user, api]'
+# Include 'user' if the change is relevant to end users.
+# Include 'api' if there is a change to a library API.
+# Default: '[user]'
+change_logs: [user]
\ No newline at end of file
diff --git a/receiver/azureeventhubreceiver/eventhubhandler.go b/receiver/azureeventhubreceiver/eventhubhandler.go
index bf48cb1047f9..e45bb8ce7063 100644
--- a/receiver/azureeventhubreceiver/eventhubhandler.go
+++ b/receiver/azureeventhubreceiver/eventhubhandler.go
@@ -8,6 +8,7 @@ import (
 
 	eventhub "github.com/Azure/azure-event-hubs-go/v3"
 	"go.opentelemetry.io/collector/component"
+	"go.opentelemetry.io/collector/extension/experimental/storage"
 	"go.opentelemetry.io/collector/receiver"
 	"go.uber.org/zap"
 
@@ -43,11 +44,12 @@ type listerHandleWrapper interface {
 }
 
 type eventhubHandler struct {
-	hub          hubWrapper
-	dataConsumer dataConsumer
-	config       *Config
-	settings     receiver.Settings
-	cancel       context.CancelFunc
+	hub           hubWrapper
+	dataConsumer  dataConsumer
+	config        *Config
+	settings      receiver.Settings
+	cancel        context.CancelFunc
+	storageClient storage.Client
 }
 
 func (h *eventhubHandler) run(ctx context.Context, host component.Host) error {
@@ -58,6 +60,7 @@ func (h *eventhubHandler) run(ctx context.Context, host component.Host) error {
 		h.settings.Logger.Debug("Error connecting to Storage", zap.Error(err))
 		return err
 	}
+	h.storageClient = storageClient
 
 	if h.hub == nil { // set manually for testing.
 		hub, newHubErr := eventhub.NewHubFromConnectionString(h.config.Connection, eventhub.HubWithOffsetPersistence(&storageCheckpointPersister{storageClient: storageClient}))
@@ -160,6 +163,13 @@ func (h *eventhubHandler) newMessageHandler(ctx context.Context, event *eventhub
 }
 
 func (h *eventhubHandler) close(ctx context.Context) error {
+	if h.storageClient != nil {
+		if err := h.storageClient.Close(ctx); err != nil {
+			return err
+		}
+		h.storageClient = nil
+	}
+
 	if h.hub != nil {
 		err := h.hub.Close(ctx)
 		if err != nil {