Properly dispose UDS for Workload Client. (#5035) (#5044)

We can do a better job of disposing UDS for Workload Client.
Azure · Jun 9, 2021 · 472cee5 · 472cee5
1 parent 86625fa
commit 472cee5
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 23 deletions.
diff --git a/edge-util/src/Microsoft.Azure.Devices.Edge.Util/uds/HttpChunkedStreamReader.cs b/edge-util/src/Microsoft.Azure.Devices.Edge.Util/uds/HttpChunkedStreamReader.cs
@@ -122,5 +122,10 @@ public override int Read(byte[] buffer, int offset, int count)
         public override void SetLength(long value) => throw new NotImplementedException();
 
         public override void Write(byte[] buffer, int offset, int count) => throw new NotImplementedException();
+
+        protected override void Dispose(bool disposing)
+        {
+            this.stream.Dispose();
+        }
     }
 }
diff --git a/edge-util/src/Microsoft.Azure.Devices.Edge.Util/uds/HttpUdsMessageHandler.cs b/edge-util/src/Microsoft.Azure.Devices.Edge.Util/uds/HttpUdsMessageHandler.cs
@@ -20,34 +20,31 @@ public HttpUdsMessageHandler(Uri providerUri)
 
         protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
         {
-            Socket socket = await this.GetConnectedSocketAsync();
-            var stream = new HttpBufferedStream(new NetworkStream(socket, true));
-
-            var serializer = new HttpRequestResponseSerializer();
-            byte[] requestBytes = serializer.SerializeRequest(request);
-
-            Events.SendRequest(request.RequestUri);
-            await stream.WriteAsync(requestBytes, 0, requestBytes.Length, cancellationToken);
-            if (request.Content != null)
+            var endpoint = new UnixDomainSocketEndPoint(this.providerUri.LocalPath);
+            using (Socket socket = new Socket(AddressFamily.Unix, SocketType.Stream, ProtocolType.Unspecified))
             {
-                await request.Content.CopyToAsync(stream);
-            }
+                Events.Connecting(this.providerUri.LocalPath);
+                await socket.ConnectAsync(endpoint);
+                Events.Connected(this.providerUri.LocalPath);
 
-            HttpResponseMessage response = await serializer.DeserializeResponse(stream, cancellationToken);
-            Events.ResponseReceived(response.StatusCode);
+                using (var stream = new HttpBufferedStream(new NetworkStream(socket, true)))
+                {
+                    var serializer = new HttpRequestResponseSerializer();
+                    byte[] requestBytes = serializer.SerializeRequest(request);
 
-            return response;
-        }
+                    Events.SendRequest(request.RequestUri);
+                    await stream.WriteAsync(requestBytes, 0, requestBytes.Length, cancellationToken);
+                    if (request.Content != null)
+                    {
+                        await request.Content.CopyToAsync(stream);
+                    }
 
-        async Task<Socket> GetConnectedSocketAsync()
-        {
-            var endpoint = new UnixDomainSocketEndPoint(this.providerUri.LocalPath);
-            var socket = new Socket(AddressFamily.Unix, SocketType.Stream, ProtocolType.Unspecified);
-            Events.Connecting(this.providerUri.LocalPath);
-            await socket.ConnectAsync(endpoint);
-            Events.Connected(this.providerUri.LocalPath);
+                    HttpResponseMessage response = await serializer.DeserializeResponse(stream, cancellationToken);
+                    Events.ResponseReceived(response.StatusCode);
 
-            return socket;
+                    return response;
+                }
+            }
         }
 
         static class Events