Remove FromAsync usage from UDPClient (#35719)

diff --git a/src/libraries/System.Net.Sockets/src/System/Net/Sockets/UDPClient.cs b/src/libraries/System.Net.Sockets/src/System/Net/Sockets/UDPClient.cs
index c7d280f..7daf3a5 100644 (file)
--- a/src/libraries/System.Net.Sockets/src/System/Net/Sockets/UDPClient.cs
+++ b/src/libraries/System.Net.Sockets/src/System/Net/Sockets/UDPClient.cs
@@ -282,11 +282,40 @@ namespace System.Net.Sockets
             }
         }
 
+        public IAsyncResult BeginSend(byte[] datagram, int bytes, AsyncCallback? requestCallback, object? state) =>
+            BeginSend(datagram, bytes, null, requestCallback, state);
+
+        public IAsyncResult BeginSend(byte[] datagram, int bytes, string? hostname, int port, AsyncCallback? requestCallback, object? state) =>
+            BeginSend(datagram, bytes, GetEndpoint(hostname, port), requestCallback, state);
+
         public IAsyncResult BeginSend(byte[] datagram, int bytes, IPEndPoint? endPoint, AsyncCallback? requestCallback, object? state)
         {
+            ValidateDatagram(datagram, bytes, endPoint);
+
+            if (endPoint is null)
+            {
+                return _clientSocket.BeginSend(datagram, 0, bytes, SocketFlags.None, requestCallback, state);
+            }
+            else
+            {
+                CheckForBroadcast(endPoint.Address);
+                return _clientSocket.BeginSendTo(datagram, 0, bytes, SocketFlags.None, endPoint, requestCallback, state);
+            }
+        }
+
+        public int EndSend(IAsyncResult asyncResult)
+        {
+            ThrowIfDisposed();
+
+            return _active ?
+                _clientSocket.EndSend(asyncResult) :
+                _clientSocket.EndSendTo(asyncResult);
+        }
+
+        private void ValidateDatagram(byte[] datagram, int bytes, IPEndPoint? endPoint)
+        {
             ThrowIfDisposed();
 
-            // Validate input parameters.
             if (datagram == null)
             {
                 throw new ArgumentNullException(nameof(datagram));
@@ -302,18 +331,9 @@ namespace System.Net.Sockets
                 // Do not allow sending packets to arbitrary host when connected.
                 throw new InvalidOperationException(SR.net_udpconnected);
             }
-
-            if (endPoint == null)
-            {
-                return _clientSocket.BeginSend(datagram, 0, bytes, SocketFlags.None, requestCallback, state);
-            }
-
-            CheckForBroadcast(endPoint.Address);
-
-            return _clientSocket.BeginSendTo(datagram, 0, bytes, SocketFlags.None, endPoint, requestCallback, state);
         }
 
-        public IAsyncResult BeginSend(byte[] datagram, int bytes, string? hostname, int port, AsyncCallback? requestCallback, object? state)
+        private IPEndPoint? GetEndpoint(string? hostname, int port)
         {
             if (_active && ((hostname != null) || (port != 0)))
             {
@@ -340,26 +360,7 @@ namespace System.Net.Sockets
                 ipEndPoint = new IPEndPoint(addresses[i], port);
             }
 
-            return BeginSend(datagram, bytes, ipEndPoint, requestCallback, state);
-        }
-
-        public IAsyncResult BeginSend(byte[] datagram, int bytes, AsyncCallback? requestCallback, object? state)
-        {
-            return BeginSend(datagram, bytes, null, requestCallback, state);
-        }
-
-        public int EndSend(IAsyncResult asyncResult)
-        {
-            ThrowIfDisposed();
-
-            if (_active)
-            {
-                return _clientSocket.EndSend(asyncResult);
-            }
-            else
-            {
-                return _clientSocket.EndSendTo(asyncResult);
-            }
+            return ipEndPoint;
         }
 
         public IAsyncResult BeginReceive(AsyncCallback? requestCallback, object? state)
@@ -595,59 +596,46 @@ namespace System.Net.Sockets
                 mcOpt);
         }
 
-        public Task<int> SendAsync(byte[] datagram, int bytes)
-        {
-            return Task<int>.Factory.FromAsync(
-                (targetDatagram, targetBytes, callback, state) => ((UdpClient)state!).BeginSend(targetDatagram, targetBytes, callback, state),
-                asyncResult => ((UdpClient)asyncResult.AsyncState!).EndSend(asyncResult),
-                datagram,
-                bytes,
-                state: this);
-        }
+        public Task<int> SendAsync(byte[] datagram, int bytes) =>
+            SendAsync(datagram, bytes, null);
+
+        public Task<int> SendAsync(byte[] datagram, int bytes, string? hostname, int port) =>
+            SendAsync(datagram, bytes, GetEndpoint(hostname, port));
 
         public Task<int> SendAsync(byte[] datagram, int bytes, IPEndPoint? endPoint)
         {
-            return Task<int>.Factory.FromAsync(
-                (targetDatagram, targetBytes, targetEndpoint, callback, state) => ((UdpClient)state!).BeginSend(targetDatagram, targetBytes, targetEndpoint, callback, state),
-                asyncResult => ((UdpClient)asyncResult.AsyncState!).EndSend(asyncResult),
-                datagram,
-                bytes,
-                endPoint,
-                state: this);
+            ValidateDatagram(datagram, bytes, endPoint);
+
+            if (endPoint is null)
+            {
+                return _clientSocket.SendAsync(new ArraySegment<byte>(datagram, 0, bytes), SocketFlags.None);
+            }
+            else
+            {
+                CheckForBroadcast(endPoint.Address);
+                return _clientSocket.SendToAsync(new ArraySegment<byte>(datagram, 0, bytes), SocketFlags.None, endPoint);
+            }
         }
 
-        public Task<int> SendAsync(byte[] datagram, int bytes, string? hostname, int port)
+        public Task<UdpReceiveResult> ReceiveAsync()
         {
-            Tuple<byte[], string?> packedArguments = Tuple.Create(datagram, hostname);
+            ThrowIfDisposed();
 
-            return Task<int>.Factory.FromAsync(
-                (targetPackedArguments, targetBytes, targetPort, callback, state) =>
-                {
-                    byte[] targetDatagram = targetPackedArguments.Item1;
-                    string? targetHostname = targetPackedArguments.Item2;
-                    var client = (UdpClient)state!;
+            return WaitAndWrap(_clientSocket.ReceiveFromAsync(
+                new ArraySegment<byte>(_buffer, 0, MaxUDPSize),
+                SocketFlags.None,
+                _family == AddressFamily.InterNetwork ? IPEndPointStatics.Any : IPEndPointStatics.IPv6Any));
 
-                    return client.BeginSend(targetDatagram, targetBytes, targetHostname, targetPort, callback, state);
-                },
-                asyncResult => ((UdpClient)asyncResult.AsyncState!).EndSend(asyncResult),
-                packedArguments,
-                bytes,
-                port,
-                state: this);
-        }
+            async Task<UdpReceiveResult> WaitAndWrap(Task<SocketReceiveFromResult> task)
+            {
+                SocketReceiveFromResult result = await task.ConfigureAwait(false);
 
-        public Task<UdpReceiveResult> ReceiveAsync()
-        {
-            return Task<UdpReceiveResult>.Factory.FromAsync(
-                (callback, state) => ((UdpClient)state!).BeginReceive(callback, state),
-                asyncResult =>
-                {
-                    var client = (UdpClient)asyncResult.AsyncState!;
-                    IPEndPoint? remoteEP = null;
-                    byte[] buffer = client.EndReceive(asyncResult, ref remoteEP);
-                    return new UdpReceiveResult(buffer, remoteEP!);
-                },
-                state: this);
+                byte[] buffer = result.ReceivedBytes < MaxUDPSize ?
+                    _buffer.AsSpan(0, result.ReceivedBytes).ToArray() :
+                    _buffer;
+
+                return new UdpReceiveResult(buffer, (IPEndPoint)result.RemoteEndPoint);
+            }
         }
 
         private void CreateClientSocket()