diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright (c) Adam Langley + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions +are met: + +1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + +3. Neither the name of the author nor the names of his contributors + may be used to endorse or promote products derived from this software + without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE CONTRIBUTORS ``AS IS'' AND ANY EXPRESS +OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE AUTHORS OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS +OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) +HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, +STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN +ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +POSSIBILITY OF SUCH DAMAGE. diff --git a/Network/Connection.hs b/Network/Connection.hs new file mode 100644 --- /dev/null +++ b/Network/Connection.hs @@ -0,0 +1,286 @@ +----------------------------------------------------------------------------- +-- | +-- Module : Network.Connection +-- Copyright : Adam Langley +-- License : BSD3-style (see LICENSE) +-- +-- Maintainer : Adam Langley +-- Stability : experimental +-- +-- Helpful functions to deal with stream-like connections +----------------------------------------------------------------------------- +module Network.Connection + ( -- * Base connections + BaseConnection(..) + , baseConnectionFromSocket + + -- * Connection functions + , Connection + , new + , newSTM + , forkWriterThread + , forkInConnection + , close + , write + , writeAtLowWater + , read + , reada + , pushBack + ) where + +import Prelude hiding (foldl, read, catch) + +import Control.Concurrent +import Control.Concurrent.STM +import Control.Exception +import Control.Monad + +import Data.Foldable (foldl) + +import qualified Data.ByteString as B +import qualified Data.Sequence as Seq + +import Network.Socket hiding (send, sendTo, recv, recvFrom) +import Network.Socket.ByteString + +-- | A BaseConnection abstracts a stream like connection. +data BaseConnection = BaseConnection { + -- | Read, at most, the given number of bytes from the connection and return + -- a ByteString of the data. EOF is signaled by an exception and a zero + -- length string is never a valid return value + baseRead :: Int -> IO B.ByteString + -- | Write the given ByteString to the connection. The write may write less + -- than the requested number of bytes (but must always write at least one + -- byte) + , baseWrite :: B.ByteString -> IO Int + -- | Close a connection + , baseClose :: IO () + } + +-- | Return a BaseConnection for the given socket. +baseConnectionFromSocket :: Socket -> BaseConnection +baseConnectionFromSocket sock = BaseConnection read write close where + read = recv sock + write = send sock + close = sClose sock + +-- | A Connection uses the functions from a BaseConnection and wraps them a +-- number of commonly needed behaviours. +-- +-- Firstly, a write queue is introduced so that writes can be non-blocking. +-- +-- Secondly, the Connection can manage a number of threads. Almost always +-- there will be a writer thread which is taking items from the write queue +-- and writing them to the BaseConnection. In addition, there can be zero or +-- more other threads managed by the Connection. If a thread which is managed +-- dies, by throwing an exception or otherwise, it will close the connection +-- and all other managed threads will be killed. +-- +-- There is also the concept of pushing data back into the Connection. This +-- is useful in a chain of reader functions where, for efficiency reasons, +-- you would want to read large blocks at a time, but the data is +-- self-deliminating so you would otherwise end up in a situation where you +-- had read too much. See the pushBack function for details. +data Connection = Connection { connbase :: BaseConnection + , connoutq :: TVar (Seq.Seq B.ByteString) + , connthreads :: TVar [ThreadId] + , connpushback :: TVar (Seq.Seq B.ByteString) + , conndeath :: IO () + , conndead :: TVar Bool } + +updateTVar :: TVar a -> (a -> a) -> STM () +updateTVar tvar f = do + v <- readTVar tvar + writeTVar tvar $ f v + +-- | Create a new Connection from a BaseConnection object +new :: IO () -- ^ the action to run when the connection closes + -> BaseConnection -- ^ the socket-like object to make a connection from + -> IO Connection +new deathaction baseconn = do + conn <- atomically $ newSTM deathaction baseconn + forkWriterThread conn + return conn + +-- | This creates most of a Connection, purely in the STM monad. The Connection +-- returned from this must be passed to forkWriterThread, otherwise nothing +-- will ever get written. +newSTM :: IO () -- ^ the action run when the connection closes + -> BaseConnection -- ^ the socket-like object to make a connection from + -> STM Connection +newSTM deathaction baseconn = do + dead <- newTVar False + outq <- newTVar Seq.empty + pushback <- newTVar Seq.empty + threads <- newTVar [] + + return $ Connection baseconn outq threads pushback deathaction dead + +-- | If you created the Connection in the STM monad using newSTM, you need to +-- call this on it in order to create the thread which processes the outgoing +-- queue. +forkWriterThread :: Connection -- ^ the connection to fork the writer thread for + -> IO () +forkWriterThread conn = do + sync <- atomically $ newTVar False + writer <- forkIO $ waitForReadySignal sync $ + connectionThreadWrapper conn $ + seqToSocket (connoutq conn) $ baseWrite $ connbase conn + -- update the thread ids in the Connection and set the ready flag + atomically (updateTVar (connthreads conn) ((:) writer) >> + writeTVar sync True) + +-- | Run the given action, as if by forkIO, and manage the thread. If the given +-- action completes or throws an exception, the connection will be closed and +-- all other managed threads will be killed +forkInConnection :: Connection -- ^ the connection to close on death + -> IO () -- ^ the action to run + -> IO () +forkInConnection conn action = do + sync <- atomically $ newTVar False + thread <- forkIO $ waitForReadySignal sync $ + connectionThreadWrapper conn action + atomically (updateTVar (connthreads conn) ((:) thread) >> + writeTVar sync True) + +-- | Wait for the given TVar to be true and then run the given action +waitForReadySignal :: (TVar Bool) -> IO a -> IO a +waitForReadySignal sync action = do + atomically (do go <- readTVar sync + if go == True then return () else retry) + action + +killThreads :: Connection -> IO () +killThreads conn = do + isDead <- atomically $ do + dead <- readTVar (conndead conn) + when (not dead) $ writeTVar (conndead conn) True + return dead + when (not isDead) $ do + t <- atomically (readTVar $ connthreads conn) + me <- myThreadId + mapM_ killThread $ filter ((/=) me) t + baseClose $ connbase conn + conndeath conn + +-- | Not all exceptions are safe to catch because of the way the GC works. If a +-- thread is killed because it's waiting on a TVar which is now garbage (e.g. +-- our writer thread when the Connection goes out of scope), all ForeignPtrs +-- held by the thread are also garbage, /at the same time/. Thus we can end +-- up holding invalid ForeignPtrs if we catch unsafe exceptions and try to +-- cleanup. +safeException :: Exception -> Maybe Exception +safeException (AsyncException _) = Nothing +safeException BlockedOnDeadMVar = Nothing +safeException BlockedIndefinitely = Nothing +safeException x = Just x + +-- | Wrap a connection thread so that, when the thread dies, it races to set +-- the dead flag. If it does so, it closes the socket and kills the other +-- threads +connectionThreadWrapper :: Connection -> IO a -> IO a +connectionThreadWrapper conn action = do + handleJust safeException (\e -> killThreads conn >> throwIO e) action + +-- | Close a connection +close :: Connection -> IO () +close = killThreads + +-- | Enqueue a ByteString to a connection. This does not block. +write :: Connection -> B.ByteString -> STM () +write conn bs = do + s <- readTVar $ connoutq conn + writeTVar (connoutq conn) (bs Seq.<| s) + +-- | Block until the write queue has less than the given number of bytes in it +-- then enqueue a new ByteString. +writeAtLowWater :: Int -- ^ the max number of bytes in the queue before we enqueue anything + -> Connection -- ^ the connection to write to + -> B.ByteString -- ^ the data to enqueue + -> STM () +writeAtLowWater lw conn bs = do + q <- readTVar $ connoutq conn + let size = foldl (\sz bs -> sz + B.length bs) 0 q + if size > lw + then retry + else writeTVar (connoutq conn) $ bs Seq.<| q + +-- | Read some number of bytes from a connection. The size is only a hint, +-- the returned data may be shorter. A zero length read is EOF +read :: Connection -> Int -> IO B.ByteString +read conn sz = do + pb <- atomically $ do + pushback <- readTVar $ connpushback conn + case Seq.viewl pushback of + Seq.EmptyL -> return Nothing + head Seq.:< rest -> + if B.length head <= sz + then do + writeTVar (connpushback conn) rest + return $ Just head + else do + let (left, right) = B.splitAt sz head + writeTVar (connpushback conn) $ right Seq.<| rest + return $ Just left + case pb of + Nothing -> (baseRead $ connbase conn) sz + Just bs -> return bs + +-- | Read exactly a give number of bytes +reada :: Connection -> Int -> IO B.ByteString +reada conn n = do + bytes <- read conn n + when (B.null bytes) $ fail "EOF in reada" + let remaining = n - B.length bytes + if remaining == 0 + then return bytes + else reada conn remaining >>= return . B.append bytes + +-- | Unread some amount of data. It will be returned in the next call to read. +-- +-- The function pushes data to the front of the queue. Thus you need to push +-- all the data base in one go, or the order of future reads will be wrong. +-- +-- This might seem like an error, but consider the case of two actions: +-- the first reads 20 bytes and pushs back the last 10 of them. The second +-- reads 5 bytes and pushs back the last 4. If we appended to the push back +-- queue the second action would put those 4 bytes after the remaining 5 from +-- the first action. +pushBack :: Connection -> B.ByteString -> STM () +pushBack conn bs + | B.null bs = return () + | otherwise = do + pushback <- readTVar $ connpushback conn + writeTVar (connpushback conn) $ bs Seq.<| pushback + +-- | Atomically take elements from the end of the given sequence and write them +-- to the given socket. Throw an exception when the write fails +seqToSocket :: TVar (Seq.Seq B.ByteString) -- ^ data is removed from the end + -> (B.ByteString -> IO Int) -- ^ the write function + -> IO () +seqToSocket q write = do + -- Atomically remove an element from the end of the sequence + bs <- atomically (do q' <- readTVar q + (bs, rest) <- + case Seq.viewr q' of + Seq.EmptyR -> retry + rest Seq.:> head -> return (head, rest) + writeTVar q rest + return bs) + -- Write the data to the socket + writea write bs + seqToSocket q write + +-- | Write a given number of bytes to a socket. This wraps a write function +-- which may write less than the requested number of bytes so that the whole +-- of the given ByteString is written out. +writea :: (B.ByteString -> IO Int) -- ^ the write function + -> B.ByteString -- ^ the data to write + -> IO () +writea write bytes + | B.null bytes = return () + | otherwise = do + n <- write bytes + if n == B.length bytes + then return () + else writea write $ B.drop n bytes diff --git a/Setup.lhs b/Setup.lhs new file mode 100644 --- /dev/null +++ b/Setup.lhs @@ -0,0 +1,3 @@ +#!/usr/bin/env runhaskell +> import Distribution.Simple +> main = defaultMain diff --git a/network-connection.cabal b/network-connection.cabal new file mode 100644 --- /dev/null +++ b/network-connection.cabal @@ -0,0 +1,16 @@ +name: network-connection +version: 0.1.1 +license: BSD3 +license-file: LICENSE +author: Adam Langley +synopsis: A wrapper around a generic stream-like connection +description: Many uses of Sockets needs a number of common functions, like a write queue, pushback etc. This provides such and allows code to be written generically for any type of connection, be a raw socket or an SSL connection. +homepage: http://darcs.imperialviolet.org/network-connection +category: Network +build-depends: base, containers, bytestring>=0.9, network-bytestring>=0.1.1.2, network>=2.1, stm>=2.1 +stability: provisional +tested-with: GHC == 6.8.2 +exposed-modules: Network.Connection +ghc-options: -Wall -fno-warn-name-shadowing +extensions: OverloadedStrings +build-type: Simple