Splice.hsc

-- |
-- Module      : Network.Socket.Splice
-- Copyright   : (c) Cetin Sert 2012
-- License     : BSD-style
--
-- Maintainer  : fusion@corsis.eu
-- Stability   : stable
-- Portability : GHC-only

#ifdef LINUX_SPLICE
#include <fcntl.h>
{-# LANGUAGE CPP #-}
{-# LANGUAGE ForeignFunctionInterface #-}
#endif

module Network.Socket.Splice (
    Length
  , zeroCopy
  , loopSplice
#ifdef LINUX_SPLICE
  , c_splice
#endif
  ) where

import Data.Word
import Foreign.Ptr

import Network.Socket
import Control.Monad
import Control.Exception
import System.IO
import System.Posix.Types
import System.Posix.Internals
import GHC.IO.Handle.FD

#ifdef LINUX_SPLICE
import Data.Int
import Data.Bits
import Unsafe.Coerce
import Foreign.C.Types
import Foreign.C.Error
import System.Posix.IO
#else
import Foreign.Marshal.Alloc
#endif


-- | Indicates whether 'loopSplice' uses zero copy system calls
--   or the portable user mode Haskell substitue implementation.
zeroCopy :: Bool -- ^ True: system calls; otherwise: portable.
zeroCopy =
#ifdef LINUX_SPLICE
  True
#else
  False
#endif


--------------------------------------------------------------------------------


type Length =
#ifdef LINUX_SPLICE
  (#type size_t)
#else
  Int
#endif

try_ :: IO () -> IO ()
try_ a = (try a :: IO (Either SomeException ())) >> return ()


-- | The 'loopSplice' function pipes data from
--   one socket to another in an infinite loop.
--   On Linux this happens in kernel space with
--   zero copying, between kernel and user space.
--   On other operating systems, a portable
--   implementation utilizes a user space buffer
--   and works on handles instead of file descriptors.
loopSplice
  :: Length    -- ^ Splice length
  -> Socket    -- ^ Source socket
  -> Socket    -- ^ Target socket
  -> IO ()
loopSplice len sIn sOut = do

  let throwRecv0 = error "Network.Socket.Splice.splice ended"

  let fdIn  = fdSocket sIn
  let fdOut = fdSocket sOut 

#ifdef LINUX_SPLICE

  print "LINUX-SPLICE"

  (r,w) <- createPipe     -- r: read  end of pipe
  print ('+',r,w)         -- w: write end of pipe
  let s = Fd fdIn         -- s: source socket
  let t = Fd fdOut        -- t: target socket
  let n = nullPtr  
  let u = unsafeCoerce :: (#type ssize_t) -> (#type size_t)
  let check = throwErrnoIfMinus1 "Network.Socket.Splice.splice"
  let flags = sPLICE_F_MOVE .|. sPLICE_F_MORE
  let setNonBlockingMode v = do setNonBlockingFD fdIn  v
                                setNonBlockingFD fdOut v

  setNonBlockingMode False
  finally
    (forever $ do 
       bytes <- check $ c_splice s n w n    len    flags
       if bytes > 0
         then           c_splice r n t n (u bytes) flags
         else           throwRecv0)
    (do closeFd r
        closeFd w
        try_ $ setNonBlockingMode True
        print ('-',r,w))

#else

  s <- fdToHandle fdIn
  t <- fdToHandle fdOut
  hSetBuffering s NoBuffering
  hSetBuffering t NoBuffering
  a <- mallocBytes len :: IO (Ptr Word8)

  print "PORTABLE-SPLICE"

  finally
    (forever $ do
       bytes <-   hGetBufSome s a len
       if bytes > 0
         then     hPutBuf     t a bytes
         else     throwRecv0)
    (do free a
        try_ $ hClose s
        try_ $ hClose t)

#endif


--------------------------------------------------------------------------------


#ifdef LINUX_SPLICE
-- SPLICE

-- fcntl.h
-- ssize_t splice(
--   int          fd_in,
--   loff_t*      off_in,
--   int          fd_out,
--   loff_t*      off_out,
--   size_t       len,
--   unsigned int flags
-- );

foreign import ccall "splice"
  c_splice
  :: Fd
  -> Ptr (#type loff_t)
  -> Fd
  -> Ptr (#type loff_t)
  -> (#type size_t)
  -> Word
  -> IO (#type ssize_t)

sPLICE_F_MOVE :: Word
sPLICE_F_MOVE = (#const "SPLICE_F_MOVE")

sPLICE_F_MORE :: Word
sPLICE_F_MORE = (#const "SPLICE_F_MORE")
#endif