diff --git a/Data/Conduit/LogFile.hs b/Data/Conduit/LogFile.hs new file mode 100644 --- /dev/null +++ b/Data/Conduit/LogFile.hs @@ -0,0 +1,141 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TemplateHaskell #-} +module Data.Conduit.LogFile + ( RotatingLog + , openRotatingLog + , addChunk + , close + , defaultMaxTotal + , dummy + ) where + +import Control.Concurrent (forkIO) +import Control.Concurrent.STM (atomically) +import Control.Concurrent.STM.TBQueue +import Control.Concurrent.STM.TVar +import Control.Exception (bracket, bracketOnError, + finally) +import Control.Monad (void, when) +import qualified Data.ByteString as S +import Data.Time (UTCTime, getCurrentTime) +import Data.Word (Word) +import System.Directory (createDirectoryIfMissing, + doesFileExist, renameFile) +import System.FilePath ((<.>), ()) +import qualified System.IO as SIO +import System.IO.Unsafe (unsafePerformIO) +import System.Mem.Weak (addFinalizer) + +data Command = AddChunk !S.ByteString + | Close + +-- | Represents a folder used for totating log files. +-- +-- Since 0.2.1 +data RotatingLog = RotatingLog !(TVar State) +-- Use a data instead of a newtype so that we can attach a finalizer. + +-- | A @RotatingLog@ which performs no logging. +-- +-- Since 0.2.1 +dummy :: RotatingLog +dummy = RotatingLog $! unsafePerformIO $! newTVarIO Closed + +data State = Closed + | Running !SIO.Handle !(TBQueue Command) + +queue :: Command -> RotatingLog -> IO () +queue cmd (RotatingLog ts) = atomically $ do + s <- readTVar ts + case s of + Closed -> return () + Running _ q -> writeTBQueue q cmd + +addChunk :: RotatingLog -> S.ByteString -> IO () +addChunk lf bs = queue (AddChunk bs) lf + +close :: RotatingLog -> IO () +close = queue Close + +-- | Create a new @RotatingLog@. +-- +-- Since 0.2.1 +openRotatingLog :: FilePath -- ^ folder to contain logs + -> Word -- ^ maximum log file size, in bytes + -> IO RotatingLog +openRotatingLog dir maxTotal = do + createDirectoryIfMissing True dir + bracketOnError (moveCurrent dir) SIO.hClose $ \handle -> do + queue <- newTBQueueIO 5 + let s = Running handle queue + ts <- newTVarIO s + void $ forkIO $ loop dir ts maxTotal + let rl = RotatingLog ts + addFinalizer rl (atomically (writeTBQueue queue Close)) + return rl + +current :: FilePath -- ^ folder containing logs + -> FilePath +current = ( "current.log") + +moveCurrent :: FilePath -- ^ folder containing logs + -> IO SIO.Handle -- ^ new handle +moveCurrent dir = do + let curr = current dir + x <- doesFileExist curr + when x $ do + now <- getCurrentTime + renameFile curr $ dir suffix now + SIO.openFile curr SIO.WriteMode + +suffix :: UTCTime -> FilePath +suffix now = + (concatMap fix $ takeWhile (/= '.') $ show now) <.> "log" + where + fix ' ' = "_" + fix c | '0' <= c && c <= '9' = [c] + fix _ = "" + +loop :: FilePath -- ^ folder containing logs + -> TVar State + -> Word -- ^ maximum total log size + -> IO () +loop dir ts maxTotal = + go 0 `finally` (closeCurrentHandle `finally` atomically (writeTVar ts Closed)) + where + closeCurrentHandle = bracket + (atomically $ do + s <- readTVar ts + case s of + Closed -> return Nothing + Running h _ -> return $! Just h) + (maybe (return ()) SIO.hClose) + (const $ return ()) + + go total = do + res <- atomically $ do + s <- readTVar ts + case s of + Closed -> return Nothing + Running handle queue -> do + cmd <- readTBQueue queue + case cmd of + Close -> return Nothing + AddChunk bs -> return $! Just (handle, queue, bs) + case res of + Nothing -> return () + Just (handle, queue, bs) -> do + let total' = total + fromIntegral (S.length bs) + S.hPut handle bs + SIO.hFlush handle + if total' > maxTotal + then do + bracket + (SIO.hClose handle >> moveCurrent dir) + (\handle' -> atomically $ writeTVar ts $ Running handle' queue) + (const $ return ()) + go 0 + else go total' + +defaultMaxTotal :: Word +defaultMaxTotal = 5 * 1024 * 1024 -- 5 MB diff --git a/Data/Conduit/Process/Unix.hs b/Data/Conduit/Process/Unix.hs --- a/Data/Conduit/Process/Unix.hs +++ b/Data/Conduit/Process/Unix.hs @@ -1,49 +1,96 @@ -{-# LANGUAGE ForeignFunctionInterface, OverloadedStrings, ScopedTypeVariables #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE ForeignFunctionInterface #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} module Data.Conduit.Process.Unix - ( forkExecuteFile + ( -- * Starting processes + forkExecuteFile , killProcess , terminateProcess , waitForProcess , ProcessStatus (..) , signalProcessHandle , signalProcessHandleGroup + -- * Process tracking + -- $processTracker + + -- ** Types + , ProcessTracker + , TrackedProcess + , ProcessTrackerException (..) + -- ** Functions + , initProcessTracker + , trackProcess + , untrackProcess + + -- * Logging + , RotatingLog + , openRotatingLog + , forkExecuteLog + + -- * Monitored process + , MonitoredProcess + , monitorProcess + , terminateMonitoredProcess ) where -import Control.Applicative ((<$>), (<*>)) -import Control.Arrow ((***)) -import Control.Concurrent (forkIO) -import Control.Exception (finally, mask, onException, handle, SomeException) -import Control.Monad (unless, void, zipWithM_, when) -import Control.Monad.Trans.Class (lift) -import Data.ByteString (ByteString, null, concat, append, singleton) -import qualified Data.ByteString.Char8 as S8 -import Data.ByteString.Unsafe (unsafePackCStringFinalizer, - unsafeUseAsCStringLen, - unsafeUseAsCString) -import Data.Conduit (Sink, Source, yield, ($$)) -import Data.Conduit.Binary (sinkHandle, sourceHandle) -import Data.Conduit.List (mapM_) -import Foreign.Marshal.Alloc (free, mallocBytes, allocaBytes) -import Foreign.Ptr (castPtr, Ptr, nullPtr) -import Foreign.Storable (sizeOf, pokeElemOff) -import Prelude (Bool (..), IO, Maybe (..), - Monad (..), flip, - fromIntegral, fst, maybe, - snd, ($), (.), (==), error, id, length, (*), - head, map, const, fmap) -import System.Posix.Types (Fd) -import System.Posix.IO.ByteString (closeFd, createPipe, - fdReadBuf, fdWriteBuf, - setFdOption, FdOption (CloseOnExec)) -import System.Posix.Process.ByteString (ProcessStatus (..), - getProcessStatus, getProcessGroupIDOf) -import System.Posix.Signals (sigKILL, signalProcess, Signal, signalProcessGroup) -import System.Posix.Types (ProcessID) -import System.IO (hClose) -import Foreign.C.Types -import Foreign.C.String -import System.Process -import System.Process.Internals +import Control.Applicative ((<$>), (<*>)) +import Control.Arrow ((***)) +import Control.Concurrent (forkIO) +import Control.Concurrent (threadDelay) +import Control.Concurrent.MVar (MVar, modifyMVar, newMVar, + readMVar, swapMVar) +import Control.Exception (Exception, SomeException, + bracketOnError, finally, + handle, mask, mask_, + onException, throwIO, try) +import Control.Monad (unless, void, when, zipWithM_) +import Control.Monad.Trans.Class (lift) +import Data.ByteString (ByteString, append, concat, + null, singleton) +import qualified Data.ByteString.Char8 as S8 +import Data.ByteString.Unsafe (unsafePackCStringFinalizer, + unsafeUseAsCString, + unsafeUseAsCStringLen) +import Data.Conduit (Sink, Source, yield, ($$)) +import Data.Conduit.Binary (sinkHandle, sourceHandle) +import Data.Conduit.List (mapM_) +import qualified Data.Conduit.List as CL +import Data.Conduit.LogFile +import Data.IORef (IORef, newIORef, readIORef, + writeIORef) +import Data.Time (getCurrentTime) +import Data.Time (diffUTCTime) +import Data.Typeable (Typeable) +import Foreign.C.String +import Foreign.C.Types +import Foreign.Marshal.Alloc (allocaBytes, free, + mallocBytes) +import Foreign.Ptr (Ptr, castPtr, nullPtr) +import Foreign.Storable (pokeElemOff, sizeOf) +import Prelude (Bool (..), Either (..), IO, + Maybe (..), Monad (..), Show, + const, error, flip, fmap, + fromIntegral, fst, head, id, + length, map, maybe, show, snd, + ($), ($!), (*), (.), (<), + (==)) +import System.IO (hClose) +import System.Posix.IO.ByteString (FdOption (CloseOnExec), + closeFd, createPipe, + fdReadBuf, fdToHandle, + fdWriteBuf, setFdOption) +import System.Posix.Process.ByteString (ProcessStatus (..), + getProcessGroupIDOf, + getProcessStatus) +import System.Posix.Signals (Signal, sigKILL, + signalProcess, + signalProcessGroup) +import System.Posix.Types (CPid (..), Fd, ProcessID) +import System.Process +import System.Process.Internals (ProcessHandle (..), + ProcessHandle__ (..), + withProcessHandle_) -- | Kill a process by sending it the KILL (9) signal. -- @@ -109,7 +156,6 @@ Nothing -> return () return ph where - ignoreExceptions = handle (\(_ :: SomeException) -> return ()) cp = CreateProcess { cmdspec = RawCommand (S8.unpack cmd) (map S8.unpack args) , cwd = S8.unpack <$> mwdir @@ -121,6 +167,9 @@ , create_group = True } +ignoreExceptions :: IO () -> IO () +ignoreExceptions = handle (\(_ :: SomeException) -> return ()) + closeOnExec :: Fd -> IO () closeOnExec fd = setFdOption fd CloseOnExec True @@ -146,3 +195,226 @@ run ptrs = allocaBytes (length ptrs * sizeOf (head ptrs)) $ \res -> zipWithM_ (pokeElemOff res) [0..] ptrs >> f res + +-- $processTracker +-- +-- Ensure that child processes are killed, regardless of how the parent process exits. +-- +-- The technique used here is: +-- +-- * Create a pipe. +-- +-- * Fork a new child process that listens on the pipe. +-- +-- * In the current process, send updates about processes that should be auto-killed. +-- +-- * When the parent process dies, listening on the pipe in the child process will get an EOF. +-- +-- * When the child process receives that EOF, it kills all processes it was told to auto-kill. +-- +-- This code was originally written for Keter, but was moved to unix-process +-- conduit in the 0.2.1 release. + +foreign import ccall unsafe "launch_process_tracker" + c_launch_process_tracker :: IO CInt + +foreign import ccall unsafe "track_process" + c_track_process :: ProcessTracker -> CPid -> CInt -> IO () + +-- | Represents the child process which handles process cleanup. +-- +-- Since 0.2.1 +newtype ProcessTracker = ProcessTracker CInt + +-- | Represents a child process which is currently being tracked by the cleanup +-- child process. +-- +-- Since 0.2.1 +data TrackedProcess = TrackedProcess !ProcessTracker !(IORef MaybePid) + +data MaybePid = NoPid | Pid !CPid + +-- | Fork off the child cleanup process. +-- +-- This will ideally only be run once for your entire application. +-- +-- Since 0.2.1 +initProcessTracker :: IO ProcessTracker +initProcessTracker = do + i <- c_launch_process_tracker + if i == -1 + then throwIO CannotLaunchProcessTracker + else return $! ProcessTracker i + +-- | Since 0.2.1 +data ProcessTrackerException = CannotLaunchProcessTracker + deriving (Show, Typeable) +instance Exception ProcessTrackerException + +-- | Begin tracking the given process. If the 'ProcessHandle' refers to a +-- closed process, no tracking will occur. If the process is closed, then it +-- will be untracked automatically. +-- +-- Note that you /must/ compile your program with @-threaded@; see +-- 'waitForProcess'. +-- +-- Since 0.2.1 +trackProcess :: ProcessTracker -> ProcessHandle -> IO TrackedProcess +trackProcess pt ph@(ProcessHandle mph) = mask_ $ do + mpid <- readMVar mph + mpid' <- case mpid of + ClosedHandle{} -> return NoPid + OpenHandle pid -> do + c_track_process pt pid 1 + return $ Pid pid + ipid <- newIORef mpid' + let tp = TrackedProcess pt ipid + case mpid' of + NoPid -> return () + Pid _ -> void $ forkIO $ do + void $ waitForProcess ph + untrackProcess tp + return $! tp + +-- | Explicitly remove the given process from the tracked process list in the +-- cleanup process. +-- +-- Since 0.2.1 +untrackProcess :: TrackedProcess -> IO () +untrackProcess (TrackedProcess pt ipid) = mask_ $ do + mpid <- readIORef ipid + case mpid of + NoPid -> return () + Pid pid -> do + c_track_process pt pid 0 + writeIORef ipid NoPid + +-- | Fork and execute a subprocess, sending stdout and stderr to the specified +-- rotating log. +-- +-- Since 0.2.1 +forkExecuteLog :: ByteString -- ^ command + -> [ByteString] -- ^ args + -> Maybe [(ByteString, ByteString)] -- ^ environment + -> Maybe ByteString -- ^ working directory + -> Maybe (Source IO ByteString) -- ^ stdin + -> RotatingLog -- ^ both stdout and stderr will be sent to this location + -> IO ProcessHandle +forkExecuteLog cmd args menv mwdir mstdin rlog = bracketOnError + setupPipe + cleanupPipes + usePipes + where + setupPipe = bracketOnError + createPipe + (\(x, y) -> closeFd x `finally` closeFd y) + (\(x, y) -> (,) <$> fdToHandle x <*> fdToHandle y) + cleanupPipes (x, y) = hClose x `finally` hClose y + + usePipes pipes@(readerH, writerH) = do + (min, _, _, ph) <- createProcess CreateProcess + { cmdspec = RawCommand (S8.unpack cmd) (map S8.unpack args) + , cwd = S8.unpack <$> mwdir + , env = map (S8.unpack *** S8.unpack) <$> menv + , std_in = maybe Inherit (const CreatePipe) mstdin + , std_out = UseHandle writerH + , std_err = UseHandle writerH + , close_fds = True + , create_group = True + } + ignoreExceptions $ addAttachMessage pipes ph + void $ forkIO $ ignoreExceptions $ + (sourceHandle readerH $$ CL.mapM_ (addChunk rlog)) `finally` hClose readerH + case (min, mstdin) of + (Just h, Just source) -> void $ forkIO $ ignoreExceptions $ + (source $$ sinkHandle h) `finally` hClose h + (Nothing, Nothing) -> return () + _ -> error $ "Invariant violated: Data.Conduit.Process.Unix.forkExecuteLog" + return ph + + addAttachMessage pipes ph = withProcessHandle_ ph $ \p_ -> do + now <- getCurrentTime + case p_ of + ClosedHandle ec -> do + addChunk rlog $ S8.concat + [ "\n\n" + , S8.pack $ show now + , ": Process immediately died with exit code " + , S8.pack $ show ec + , "\n\n" + ] + cleanupPipes pipes + OpenHandle h -> do + addChunk rlog $ S8.concat + [ "\n\n" + , S8.pack $ show now + , ": Attached new process " + , S8.pack $ show h + , "\n\n" + ] + return p_ + +data Status = NeedsRestart | NoRestart | Running ProcessHandle + +-- | Run the given command, restarting if the process dies. +monitorProcess + :: (ByteString -> IO ()) -- ^ log + -> ProcessTracker + -> Maybe S8.ByteString -- ^ setuid + -> S8.ByteString -- ^ executable + -> S8.ByteString -- ^ working directory + -> [S8.ByteString] -- ^ command line parameter + -> [(S8.ByteString, S8.ByteString)] -- ^ environment + -> RotatingLog + -> IO MonitoredProcess +monitorProcess log processTracker msetuid exec dir args env rlog = do + mstatus <- newMVar NeedsRestart + let loop mlast = do + next <- modifyMVar mstatus $ \status -> + case status of + NoRestart -> return (NoRestart, return ()) + _ -> do + now <- getCurrentTime + case mlast of + Just last | diffUTCTime now last < 5 -> do + log $ "Process restarting too quickly, waiting before trying again: " `S8.append` exec + threadDelay $ 5 * 1000 * 1000 + _ -> return () + let (cmd, args') = + case msetuid of + Nothing -> (exec, args) + Just setuid -> ("sudo", "-E" : "-u" : setuid : "--" : exec : args) + res <- try $ forkExecuteLog + cmd + args' + (Just env) + (Just dir) + (Just $ return ()) + rlog + case res of + Left e -> do + log $ "Data.Conduit.Process.Unix.monitorProcess: " `S8.append` S8.pack (show (e :: SomeException)) + return (NeedsRestart, return ()) + Right pid -> do + log $ "Process created: " `S8.append` exec + return (Running pid, do + void $ trackProcess processTracker pid + void $ waitForProcess pid + loop (Just now)) + next + forkIO $ loop Nothing + return $ MonitoredProcess mstatus + +-- | Abstract type containing information on a process which will be restarted. +newtype MonitoredProcess = MonitoredProcess (MVar Status) + +-- | Terminate the process and prevent it from being restarted. +terminateMonitoredProcess :: MonitoredProcess -> IO () +terminateMonitoredProcess (MonitoredProcess mstatus) = do + status <- swapMVar mstatus NoRestart + case status of + Running pid -> do + terminateProcess pid + threadDelay 1000000 + killProcess pid + _ -> return () diff --git a/LICENSE b/LICENSE --- a/LICENSE +++ b/LICENSE @@ -1,30 +1,20 @@ -Copyright (c)2011, Michael Snoyman - -All rights reserved. - -Redistribution and use in source and binary forms, with or without -modification, are permitted provided that the following conditions are met: - - * Redistributions of source code must retain the above copyright - notice, this list of conditions and the following disclaimer. +Copyright (c) 2012 Michael Snoyman, http://www.yesodweb.com/ - * 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. +Permission is hereby granted, free of charge, to any person obtaining +a copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: - * Neither the name of Michael Snoyman nor the names of other - contributors may be used to endorse or promote products derived - from this software without specific prior written permission. +The above copyright notice and this permission notice shall be +included in all copies or substantial portions of the Software. -THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND 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 COPYRIGHT -OWNER 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. +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE +LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION +WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/cbits/process-tracker.c b/cbits/process-tracker.c new file mode 100644 --- /dev/null +++ b/cbits/process-tracker.c @@ -0,0 +1,100 @@ +#include +#include +#include +#include + +struct node { + pid_t pid; + struct node *next; +}; + +static struct node * add_node(pid_t pid, struct node *n) { + struct node *n2 = malloc(sizeof(struct node)); + n2->pid = pid; + n2->next = n; + return n2; +} + +static struct node * remove_node(pid_t pid, struct node *n) { + if (!n) { + return n; + } + else if (n->pid == pid) { + struct node *n2 = n->next; + free(n); + return remove_node(pid, n2); + } + else { + n->next = remove_node(pid, n->next); + return n; + } +} + +extern void track_process(int fd, pid_t pid, int b) { + unsigned int buffer[2]; + + //printf("Tracking process %d, %d\n", pid, b); + + buffer[0] = pid; + buffer[1] = b; + if (! write(fd, buffer, sizeof(unsigned int) * 2)) { + //printf("Error writing to fd %d\n", fd); + } +} + +// Returns FD to write to, or -1 on failure. +extern int launch_process_tracker(void) { + int pipes[2]; + pid_t child; + + if (pipe(pipes)) { + return -1; + } + + child = fork(); + + if (child < 0) { + return -1; + } + else if (child == 0) { + unsigned int buffer[2]; + struct node *n = 0, *n2; + + close(pipes[1]); + + // Prevent monitoring programs like Upstart from killing this + // new process along with the parent + setpgid(0, 0); + + while (read(pipes[0], buffer, sizeof(unsigned int) * 2) > 0) { + if (buffer[1]) { + //printf("Adding node %d\n", buffer[0]); + n = add_node(buffer[0], n); + } + else { + //printf("Removing node %d\n", buffer[0]); + n = remove_node(buffer[0], n); + } + } + + for (n2 = n; n2; n2 = n2->next) { + //printf("Sending process %d TERM signal\n", n2->pid); + kill(n2->pid, SIGTERM); + } + + sleep(2); + + while (n) { + //printf("Sending process %d KILL signal\n", n->pid); + kill(n2->pid, SIGKILL); + + n2 = n; + n = n->next; + free(n2); + } + } + else { + close(pipes[0]); + return pipes[1]; + } +} diff --git a/unix-process-conduit.cabal b/unix-process-conduit.cabal --- a/unix-process-conduit.cabal +++ b/unix-process-conduit.cabal @@ -2,7 +2,7 @@ -- documentation, see http://haskell.org/cabal/users-guide/ name: unix-process-conduit -version: 0.2.0.2 +version: 0.2.1 synopsis: Run processes on Unix systems, with a conduit interface description: This library allows you to provide @conduit@ datatypes for the input and output streams. Note that you must compile your programs with @-threaded@. homepage: https://github.com/snoyberg/conduit @@ -16,12 +16,18 @@ library exposed-modules: Data.Conduit.Process.Unix + Data.Conduit.LogFile build-depends: base >= 4 && < 5 , transformers , bytestring , conduit >= 0.5 && < 1.1 , unix >= 2.5 , process >= 1.1 + , time + , filepath + , stm + , directory + c-sources: cbits/process-tracker.c test-suite test hs-source-dirs: test