git.iko.soy/distributed-process
on fcfc19d help sign in
distributed-process/distributed-process-simplelocalnet/src/Control/Distributed/Process/Backend/SimpleLocalnet/Internal/Multicast.hs
126 lines 4.7 kB history raw
1-- | Multicast utilities
2{-# OPTIONS_GHC -fno-warn-orphans #-}
3module Control.Distributed.Process.Backend.SimpleLocalnet.Internal.Multicast (initMulticast) where
5import Data.Function (on)
6import Data.Map (Map)
7import qualified Data.Map as Map (empty)
8import Data.Binary (Binary, decode, encode)
9import Data.IORef (IORef, newIORef, readIORef, modifyIORef)
10import qualified Data.ByteString as BSS (ByteString, concat)
11import qualified Data.ByteString.Lazy as BSL
12 ( ByteString
13 , empty
14 , append
15 , fromChunks
16 , toChunks
17 , length
18 , splitAt
19 )
20import Data.Accessor (Accessor, (^:), (^.), (^=))
21import qualified Data.Accessor.Container as DAC (mapDefault)
22import Control.Applicative ((<$>))
23import Network.Socket (HostName, PortNumber, Socket, SockAddr)
24import qualified Network.Socket.ByteString as NBS (recvFrom, sendManyTo)
25import Network.Transport.Internal (decodeInt32, encodeInt32)
26import Network.Multicast (multicastSender, multicastReceiver)
28--------------------------------------------------------------------------------
29-- Top-level API --
30--------------------------------------------------------------------------------
32-- | Given a hostname and a port number, initialize the multicast system.
33--
34-- Note: it is important that you never send messages larger than the maximum
35-- message size; if you do, all subsequent communication will probably fail.
36--
37-- Returns a reader and a writer.
38--
39-- NOTE: By rights the two functions should be "locally" polymorphic in 'a',
40-- but this requires impredicative types.
41initMulticast :: forall a. Binary a
42 => HostName -- ^ Multicast IP
43 -> PortNumber -- ^ Port number
44 -> Int -- ^ Maximum message size
45 -> IO (IO (a, SockAddr), a -> IO ())
46initMulticast host port bufferSize = do
47 (sendSock, sendAddr) <- multicastSender host port
48 readSock <- multicastReceiver host port
49 st <- newIORef Map.empty
50 return (recvBinary readSock st bufferSize, writer sendSock sendAddr)
51 where
52 writer :: forall a. Binary a => Socket -> SockAddr -> a -> IO ()
53 writer sock addr val = do
54 let bytes = encode val
55 len = encodeInt32 (BSL.length bytes)
56 NBS.sendManyTo sock (len : BSL.toChunks bytes) addr
58--------------------------------------------------------------------------------
59-- UDP multicast read, dealing with multiple senders --
60--------------------------------------------------------------------------------
62type UDPState = Map SockAddr BSL.ByteString
64-- TODO: This is inefficient and an orphan instance.
65-- Requested official instance (https://github.com/haskell/network/issues/38)
66instance Ord SockAddr where
67 compare = compare `on` show
69bufferFor :: SockAddr -> Accessor UDPState BSL.ByteString
70bufferFor = DAC.mapDefault BSL.empty
72bufferAppend :: SockAddr -> BSS.ByteString -> UDPState -> UDPState
73bufferAppend addr bytes =
74 bufferFor addr ^: flip BSL.append (BSL.fromChunks [bytes])
76recvBinary :: Binary a => Socket -> IORef UDPState -> Int -> IO (a, SockAddr)
77recvBinary sock st bufferSize = do
78 (bytes, addr) <- recvWithLength sock st bufferSize
79 return (decode bytes, addr)
81recvWithLength :: Socket
82 -> IORef UDPState
83 -> Int
84 -> IO (BSL.ByteString, SockAddr)
85recvWithLength sock st bufferSize = do
86 (len, addr) <- recvExact sock 4 st bufferSize
87 let n = decodeInt32 . BSS.concat . BSL.toChunks $ len
88 bytes <- recvExactFrom addr sock n st bufferSize
89 return (bytes, addr)
91-- Receive all bytes currently in the buffer
92recvAll :: Socket -> IORef UDPState -> Int -> IO SockAddr
93recvAll sock st bufferSize = do
94 (bytes, addr) <- NBS.recvFrom sock bufferSize
95 modifyIORef st $ bufferAppend addr bytes
96 return addr
98recvExact :: Socket
99 -> Int
100 -> IORef UDPState
101 -> Int
102 -> IO (BSL.ByteString, SockAddr)
103recvExact sock n st bufferSize = do
104 addr <- recvAll sock st bufferSize
105 bytes <- recvExactFrom addr sock n st bufferSize
106 return (bytes, addr)
108recvExactFrom :: SockAddr
109 -> Socket
110 -> Int
111 -> IORef UDPState
112 -> Int
113 -> IO BSL.ByteString
114recvExactFrom addr sock n st bufferSize = go
115 where
116 go :: IO BSL.ByteString
117 go = do
118 accAddr <- (^. bufferFor addr) <$> readIORef st
119 if BSL.length accAddr >= fromIntegral n
120 then do
121 let (bytes, accAddr') = BSL.splitAt (fromIntegral n) accAddr
122 modifyIORef st $ bufferFor addr ^= accAddr'
123 return bytes
124 else do
125 _ <- recvAll sock st bufferSize
126 go