git.iko.soy/distributed-process
on 3acfb21 help sign in
distributed-process/distributed-process-simplelocalnet/src/Control/Distributed/Process/Backend/SimpleLocalnet/Internal/Multicast.hs
128 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#if MIN_VERSION_network(2,4,0)
65-- network-2.4.0 provides the Ord instance for us
66#else
67instance Ord SockAddr where
68 compare = compare `on` show
69#endif
71bufferFor :: SockAddr -> Accessor UDPState BSL.ByteString
72bufferFor = DAC.mapDefault BSL.empty
74bufferAppend :: SockAddr -> BSS.ByteString -> UDPState -> UDPState
75bufferAppend addr bytes =
76 bufferFor addr ^: flip BSL.append (BSL.fromChunks [bytes])
78recvBinary :: Binary a => Socket -> IORef UDPState -> Int -> IO (a, SockAddr)
79recvBinary sock st bufferSize = do
80 (bytes, addr) <- recvWithLength sock st bufferSize
81 return (decode bytes, addr)
83recvWithLength :: Socket
84 -> IORef UDPState
85 -> Int
86 -> IO (BSL.ByteString, SockAddr)
87recvWithLength sock st bufferSize = do
88 (len, addr) <- recvExact sock 4 st bufferSize
89 let n = decodeInt32 . BSS.concat . BSL.toChunks $ len
90 bytes <- recvExactFrom addr sock n st bufferSize
91 return (bytes, addr)
93-- Receive all bytes currently in the buffer
94recvAll :: Socket -> IORef UDPState -> Int -> IO SockAddr
95recvAll sock st bufferSize = do
96 (bytes, addr) <- NBS.recvFrom sock bufferSize
97 modifyIORef st $ bufferAppend addr bytes
98 return addr
100recvExact :: Socket
101 -> Int
102 -> IORef UDPState
103 -> Int
104 -> IO (BSL.ByteString, SockAddr)
105recvExact sock n st bufferSize = do
106 addr <- recvAll sock st bufferSize
107 bytes <- recvExactFrom addr sock n st bufferSize
108 return (bytes, addr)
110recvExactFrom :: SockAddr
111 -> Socket
112 -> Int
113 -> IORef UDPState
114 -> Int
115 -> IO BSL.ByteString
116recvExactFrom addr sock n st bufferSize = go
117 where
118 go :: IO BSL.ByteString
119 go = do
120 accAddr <- (^. bufferFor addr) <$> readIORef st
121 if BSL.length accAddr >= fromIntegral n
122 then do
123 let (bytes, accAddr') = BSL.splitAt (fromIntegral n) accAddr
124 modifyIORef st $ bufferFor addr ^= accAddr'
125 return bytes
126 else do
127 _ <- recvAll sock st bufferSize
128 go