Deliver a receive through STM, and stop pretending readiness works on Windows

CHANGELOG.md

@@ -17,6 +17,10 @@
   [#626](https://github.com/haskell/network/pull/626)
 * WINIO: fixing a spurious EOF when a recv completes synchronously.
   [#626](https://github.com/haskell/network/pull/626)
+* Adding recvSTM, recvFromSTM, recvBufSTM and recvBufFromSTM, which
+  deliver a receive through STM and work on Windows too.  The readiness
+  based waitReadSocketSTM and waitWriteSocketSTM cannot work there and
+  now say so instead of blocking forever.
 
 ## Version 3.2.9.0
 

Network/Socket/ByteString.hs

@@ -32,6 +32,8 @@ module Network.Socket.ByteString (
     -- * Receive data from a socket
     recv,
     recvFrom,
+    recvSTM,
+    recvFromSTM,
 
     -- * Advanced send and recv
     sendMsg,

Network/Socket/ByteString/IO.hsc

@@ -29,6 +29,8 @@ module Network.Socket.ByteString.IO
     -- * Receive data from a socket
     , recv
     , recvFrom
+    , recvSTM
+    , recvFromSTM
     , waitWhen0
 
     -- * Advanced send and recv
@@ -58,6 +60,8 @@ import System.Posix.Types (Fd(..))
 
 import Network.Socket.Flag
 import Network.Socket.SockAddr (annotateWithSocket)
+import Network.Socket.STM (viaSTM)
+import Control.Concurrent.STM (STM)
 #if defined(mingw32_HOST_OS)
 import Network.Socket.Win32.Load
 #endif
@@ -356,3 +360,20 @@ recvMsg s siz clen flags = recvMsg' `annotateWithSocket` (s, Nothing)
             let bs' | len < siz = PS fptr 0 len
                     | otherwise = bs
             return (addr, bs', cmsgs, flags')
+
+-- | Start receiving and return an 'STM' action delivering the data,
+--   together with an action cancelling the receive.  Unlike
+--   'Network.Socket.waitReadSocketSTM' this works on Windows too,
+--   because it waits for the receive to complete rather than for the
+--   socket to become readable.
+--
+--   Cancelling can lose data: the receive may already have taken a
+--   datagram out of the kernel queue.  Only abandon the 'STM' action
+--   where losing it is acceptable, such as on shutdown.
+recvSTM :: Socket -> Int -> IO (STM ByteString, IO ())
+recvSTM s nbytes = viaSTM $ recv s nbytes
+
+-- | 'recvSTM' for unconnected sockets, also returning the peer address.
+--   The same caveat about cancelling applies.
+recvFromSTM :: Socket -> Int -> IO (STM (ByteString, SockAddr), IO ())
+recvFromSTM s nbytes = viaSTM $ recvFrom s nbytes

Network/Socket/STM.hs

@@ -1,24 +1,139 @@
-module Network.Socket.STM where
+{-# LANGUAGE CPP #-}
+{-# LANGUAGE LambdaCase #-}
 
-import Control.Concurrent
+-- | STM interfaces for sockets.
+--
+-- Two styles are offered here:
+--
+-- * Readiness: 'waitReadSocketSTM' and friends return an 'STM' action
+--   that becomes available once the socket is ready.  This mirrors
+--   @select@\/@epoll@ and is __POSIX only__: Windows completion ports
+--   report that an operation has finished, not that one could be
+--   started, so there is no readiness to wait for.  All four throw on
+--   Windows.
+--
+-- * Completion: 'recvBufSTM' and 'recvBufFromSTM' start a receive and
+--   return an 'STM' action that delivers its result.  These work on
+--   every platform and are what you want if the code has to run on
+--   Windows.  "Network.Socket.ByteString" has 'ByteString' versions.
+module Network.Socket.STM (
+    -- * Waiting for readiness (POSIX only)
+    waitReadSocketSTM,
+    waitAndCancelReadSocketSTM,
+    waitWriteSocketSTM,
+    waitAndCancelWriteSocketSTM,
+
+    -- * Receiving through STM (all platforms)
+    recvBufSTM,
+    recvBufFromSTM,
+
+    -- * Building block
+    viaSTM,
+) where
+
+import Control.Concurrent (forkIO, killThread)
 import Control.Concurrent.STM
+import qualified Control.Exception as E
+import Network.Socket.Buffer
+import Network.Socket.Imports
 import Network.Socket.Types
-import System.Posix.Types
+#if !defined(mingw32_HOST_OS)
+import Control.Concurrent (threadWaitReadSTM, threadWaitWriteSTM)
+import System.Posix.Types (Fd (..))
+#endif
 
 -- | STM action to wait until the socket is ready for reading.
+--
+--   __POSIX only.__  On Windows this throws: completion ports have no
+--   notion of readiness.  Use 'recvBufFromSTM', or
+--   @Network.Socket.ByteString.recvFromSTM@, instead.
 waitReadSocketSTM :: Socket -> IO (STM ())
 waitReadSocketSTM s = fst <$> waitAndCancelReadSocketSTM s
 
 -- | STM action to wait until the socket is ready for reading and STM
 --   action to cancel the waiting.
+--
+--   __POSIX only.__  See 'waitReadSocketSTM'.
 waitAndCancelReadSocketSTM :: Socket -> IO (STM (), IO ())
+#if defined(mingw32_HOST_OS)
+waitAndCancelReadSocketSTM _ =
+    ioError $
+        userError $
+            "waitAndCancelReadSocketSTM: Windows completion ports do not "
+                ++ "provide readiness notification; use recvBufFromSTM or "
+                ++ "Network.Socket.ByteString.recvFromSTM instead"
+#else
 waitAndCancelReadSocketSTM s = withFdSocket s $ threadWaitReadSTM . Fd . fromIntegral
+#endif
 
 -- | STM action to wait until the socket is ready for writing.
+--
+--   __POSIX only.__  On Windows this throws.  Completion ports give no
+--   way to ask whether a send would block, and none is needed: issue
+--   the send and let it complete asynchronously.
 waitWriteSocketSTM :: Socket -> IO (STM ())
 waitWriteSocketSTM s = fst <$> waitAndCancelWriteSocketSTM s
 
 -- | STM action to wait until the socket is ready for writing and STM
 --   action to cancel the waiting.
+--
+--   __POSIX only.__  See 'waitWriteSocketSTM'.
 waitAndCancelWriteSocketSTM :: Socket -> IO (STM (), IO ())
+#if defined(mingw32_HOST_OS)
+waitAndCancelWriteSocketSTM _ =
+    ioError $
+        userError $
+            "waitAndCancelWriteSocketSTM: Windows completion ports do not "
+                ++ "provide readiness notification, and none is needed for "
+                ++ "sending: issue the send instead"
+#else
 waitAndCancelWriteSocketSTM s = withFdSocket s $ threadWaitWriteSTM . Fd . fromIntegral
+#endif
+
+-- | Start receiving into the given buffer and return an 'STM' action
+--   delivering the number of bytes received, together with an action
+--   cancelling the receive.  If the receive fails, the 'STM' action
+--   rethrows the exception.
+--
+--   Unlike 'waitReadSocketSTM' this works on Windows, because it waits
+--   for a completion rather than for readiness.
+--
+--   Two consequences follow from that, and both matter when composing
+--   with 'orElse':
+--
+--   * Cancelling can lose data.  The receive may already have taken a
+--     datagram out of the kernel queue, and that datagram is then gone.
+--     Only abandon the 'STM' action on paths where losing it is
+--     acceptable, such as shutdown.
+--
+--   * Cancelling needs the receive to be interruptible.  It is on POSIX
+--     and under WinIO, but not under the old Windows I/O manager, where
+--     the receive blocks in a foreign call and the cancel action waits
+--     for it.
+--
+--   * The buffer must stay alive until the 'STM' action completes or
+--     the cancel action returns.
+recvBufSTM :: Socket -> Ptr Word8 -> Int -> IO (STM Int, IO ())
+recvBufSTM s ptr nbytes = viaSTM $ recvBuf s ptr nbytes
+
+-- | 'recvBufSTM' for unconnected sockets, also returning the peer
+--   address.  The same caveats apply.
+recvBufFromSTM
+    :: SocketAddress sa => Socket -> Ptr Word8 -> Int -> IO (STM (Int, sa), IO ())
+recvBufFromSTM s ptr nbytes = viaSTM $ recvBufFrom s ptr nbytes
+
+-- | Run a blocking socket operation in a separate thread and hand its
+--   result over through STM, together with an action cancelling it.
+--   This is what 'recvBufSTM' is built from; the same caveats about
+--   cancelling and about buffer lifetime apply to anything built with
+--   it.
+viaSTM :: IO a -> IO (STM a, IO ())
+viaSTM act = do
+    var <- newTVarIO Nothing
+    tid <- forkIO $ E.try act >>= atomically . writeTVar var . Just
+    let wait =
+            readTVar var >>= \case
+                Nothing -> retry
+                Just (Left e) -> throwSTM (e :: E.SomeException)
+                Just (Right x) -> return x
+    return (wait, killThread tid)

network.cabal

@@ -204,6 +204,7 @@ test-suite spec
         directory,
         HUnit,
         network,
+        stm,
         temporary,
         hspec >=2.6,
         QuickCheck

tests/Network/Socket/ByteStringSpec.hs

@@ -1,8 +1,11 @@
+{-# LANGUAGE CPP #-}
 {-# LANGUAGE OverloadedStrings #-}
 
 module Network.Socket.ByteStringSpec (main, spec) where
 
+import Control.Concurrent (threadDelay)
 import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
+import Control.Concurrent.STM (atomically)
 import Control.Monad
 import Data.Bits
 import qualified Data.ByteString as S
@@ -21,6 +24,40 @@
 
 spec :: Spec
 spec = do
+    describe "recvSTM" $ do
+        it "delivers the data through STM" $ do
+            let server sock = do
+                    (wait, _cancel) <- recvSTM sock 1024
+                    atomically wait `shouldReturn` testMsg
+                client sock = send sock testMsg
+            tcpTest client server
+
+#if !defined(mingw32_HOST_OS)
+        -- Cancelling needs the underlying receive to be interruptible.
+        -- It is on POSIX and under WinIO, but not under the old Windows
+        -- I/O manager, where the receive blocks in a foreign call that
+        -- killThread cannot reach.
+        it "can be cancelled" $ do
+            let server sock = do
+                    (_wait, cancel) <- recvSTM sock 1024
+                    cancel
+                    -- the socket is still usable afterwards
+                    recv sock 1024 `shouldReturn` testMsg
+                client sock = do
+                    threadDelay 100000
+                    void $ send sock testMsg
+            tcpTest client server
+#endif
+
+    describe "recvFromSTM" $ do
+        it "delivers the datagram and the peer address through STM" $ do
+            let server sock = do
+                    (wait, _cancel) <- recvFromSTM sock 1024
+                    (bs, _) <- atomically wait
+                    bs `shouldBe` testMsg
+                client sock serverAddr' = void $ sendTo sock testMsg serverAddr'
+            udpTest client server
+
     describe "send" $ do
         it "works well" $ do
             let server sock = recv sock 1024 `shouldReturn` testMsg