Linux.hsc

{-# LANGUAGE CApiFFI #-}
-- Let sendfile() take a 64-bit off_t also on 32-bit platforms.
{-# OPTIONS_GHC -optc-D_FILE_OFFSET_BITS=64 #-}

#define _FILE_OFFSET_BITS 64

module Network.Sendfile.Linux (
    sendfile
  , sendfile'
  , sendfileFd
  , sendfileFd'
  , sendfileWithHeader
  , sendfileFdWithHeader
  ) where

import Control.Exception
import Control.Monad
import Data.ByteString (ByteString)
import qualified Data.ByteString as B
import Data.ByteString.Unsafe (unsafeUseAsCStringLen)
import Foreign.C.Error (eAGAIN, eINTR, getErrno, throwErrno)
import Foreign.C.Types
import Foreign.Marshal (alloca)
import Foreign.Ptr (Ptr, plusPtr)
import Foreign.Storable (poke, sizeOf)
import GHC.Conc (threadWaitWrite)
import Network.Sendfile.Types
import Network.Socket
import System.Posix.Files
import System.Posix.IO ( OpenMode(..)
                       , OpenFileFlags(..)
                       , defaultFileFlags
                       , closeFd
                       )
import System.Posix.Types

#include <sys/sendfile.h>
#include <sys/socket.h>

isLargeSize :: Bool
isLargeSize = sizeOf (0 :: CSize) == 8

safeSize :: CSize
safeSize
  | isLargeSize = 2^(60 :: Int)
  | otherwise   = 2^(30 :: Int)

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

-- |
-- Simple binding for sendfile() of Linux.
-- Used system calls:
--
--  - EntireFile -- open(), stat(), sendfile(), and close()
--
--  - PartOfFile -- open(), sendfile(), and close()
--
-- If the size of the file is unknown when sending the entire file,
-- specifying PartOfFile is much faster.
--
-- The fourth action argument is called when a file is sent as chunks.
-- Chucking is inevitable if the socket is non-blocking (this is the
-- default) and the file is large. The action is called after a chunk
-- is sent and before waiting the socket to be ready for writing.

sendfile :: Socket -> FilePath -> FileRange -> IO () -> IO ()
sendfile sock path range hook = bracket setup teardown $ \fd ->
    sendfileFd sock fd range hook
  where
    setup = openFd path ReadOnly defaultFileFlags{nonBlock=True}
    teardown = closeFd

sendfile' :: Fd -> ByteString -> FileRange -> IO () -> IO ()
sendfile' dst path range hook = bracket setup teardown $ \src ->
    sendfileFd' dst src range hook
  where
    setup = openFdBS path ReadOnly defaultFileFlags{nonBlock=True}
    teardown = closeFd

-- |
-- Simple binding for sendfile() of Linux.
-- Used system calls:
--
--  - EntireFile -- stat() and sendfile()
--
--  - PartOfFile -- sendfile()
--
-- If the size of the file is unknown when sending the entire file,
-- specifying PartOfFile is much faster.
--
-- The fourth action argument is called when a file is sent as chunks.
-- Chucking is inevitable if the socket is non-blocking (this is the
-- default) and the file is large. The action is called after a chunk
-- is sent and before waiting the socket to be ready for writing.
sendfileFd :: Socket -> Fd -> FileRange -> IO () -> IO ()
sendfileFd sock fd range hook = do
  withFdSocket sock $ \s -> do
    let dst = Fd s
    sendfileFd' dst fd range hook

sendfileFd' :: Fd -> Fd -> FileRange -> IO () -> IO ()
sendfileFd' dst src range hook = do
  checkRange range
  alloca $ \offp -> case range of
        EntireFile -> do
            poke offp 0
            -- System call is very slow. Use PartOfFile instead.
            len <- fileSize <$> getFdStatus src
            sendfileloop dst src offp len hook
        PartOfFile off len -> do
            poke offp (fromIntegral off)
            sendfileloop dst src offp (fromIntegral len) hook

-- The remaining length is COff since CSize is 32-bit on 32-bit platforms.
sendfileloop :: Fd -> Fd -> Ptr COff -> COff -> IO () -> IO ()
sendfileloop dst src offp len hook = do
    -- Multicore IO manager use edge-trigger mode.
    -- So, calling threadWaitWrite only when errnor is eAGAIN.
    let toSend
          | len > fromIntegral safeSize = safeSize
          | otherwise                   = fromIntegral len
    bytes <- c_sendfile dst src offp toSend
    case bytes of
        -1 -> do
            errno <- getErrno
            if errno == eAGAIN then do
                threadWaitWrite dst
                sendfileloop dst src offp len hook
              else if errno == eINTR then
                sendfileloop dst src offp len hook
              else
                throwErrno "Network.Sendfile.Linux.sendfileloop"
        0  -> return () -- the file is truncated
        _  -> do
            hook
            let left = len - fromIntegral bytes
            when (left /= 0) $ sendfileloop dst src offp left hook

-- Dst Src in order. take care
foreign import capi unsafe "sys/sendfile.h sendfile"
    c_sendfile :: Fd -> Fd -> Ptr COff -> CSize -> IO CSsize

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

-- |
-- Simple binding for send() and sendfile() of Linux.
-- Used system calls:
--
--  - EntireFile -- open(), send(), stat(), sendfile(), and close()
--
--  - PartOfFile -- open(), send(), sendfile(), and close()
--
-- The fifth header is sent with send() + the MSG_MORE flag. If the
-- file is small enough, the header and the file is send in a single
-- TCP packet.
--
-- If the size of the file is unknown when sending the entire file,
-- specifying PartOfFile is much faster.
--
-- The fourth action argument is called when a file is sent as chunks.
-- Chucking is inevitable if the socket is non-blocking (this is the
-- default) and the file is large. The action is called after a chunk
-- is sent and before waiting the socket to be ready for writing.

sendfileWithHeader :: Socket -> FilePath -> FileRange -> IO () -> [ByteString] -> IO ()
sendfileWithHeader sock path range hook hdr = bracket setup teardown $ \fd ->
    sendfileFdWithHeader sock fd range hook hdr
  where
    setup = openFd path ReadOnly defaultFileFlags{nonBlock=True}
    teardown = closeFd

-- |
-- Simple binding for send() and sendfile() of Linux.
-- Used system calls:
--
--  - EntireFile -- send(), stat() and sendfile()
--
--  - PartOfFile -- send() and sendfile()
--
-- The fifth header is sent with send() + the MSG_MORE flag. If the
-- file is small enough, the header and the file is send in a single
-- TCP packet.
--
-- If the size of the file is unknown when sending the entire file,
-- specifying PartOfFile is much faster.
--
-- The fourth action argument is called when a file is sent as chunks.
-- Chucking is inevitable if the socket is non-blocking (this is the
-- default) and the file is large. The action is called after a chunk
-- is sent and before waiting the socket to be ready for writing.

sendfileFdWithHeader :: Socket -> Fd -> FileRange -> IO () -> [ByteString] -> IO ()
sendfileFdWithHeader sock fd range hook hdr = do
    checkRange range
    -- Copying is much faster than syscall.
    sendMsgMore sock $ B.concat hdr
    sendfileFd sock fd range hook

sendMsgMore :: Socket -> ByteString -> IO ()
sendMsgMore sock bs = unsafeUseAsCStringLen bs $ \(buf, len) ->
    withFdSocket sock $ \fd -> sendloop (Fd fd) buf (fromIntegral len)

sendloop :: Fd -> Ptr CChar -> CSize -> IO ()
sendloop s buf len = do
    bytes <- c_send s buf len (#const MSG_MORE)
    if bytes == -1 then do
        errno <- getErrno
        if errno == eAGAIN then do
            threadWaitWrite s
            sendloop s buf len
          else if errno == eINTR then
            sendloop s buf len
          else
            throwErrno "Network.Sendfile.Linux.sendloop"
      else do
        let sent = fromIntegral bytes
        when (sent /= len) $ do
            let left = len - sent
                ptr = buf `plusPtr` fromIntegral bytes
            sendloop s ptr left

foreign import ccall unsafe "send"
  c_send :: Fd -> Ptr CChar -> CSize -> CInt -> IO CSsize