diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright Alexander Thiemann (c) 2017 + +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. + + * 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. + + * Neither the name of Alexander Thiemann nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +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. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,5 @@ +# fileplow + +[![CircleCI](https://circleci.com/gh/agrafix/fileplow.svg?style=svg)](https://circleci.com/gh/agrafix/fileplow) + +A Haskell library to process and search large files or a collection of files. This is useful if you are processing (rotated) log files for example. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/fileplow.cabal b/fileplow.cabal new file mode 100644 --- /dev/null +++ b/fileplow.cabal @@ -0,0 +1,43 @@ +name: fileplow +version: 0.1.0.0 +synopsis: Library to process and search large files or a collection of files +description: Library to process and search large files or a collection of files +homepage: https://github.com/agrafix/fileplow#readme +license: BSD3 +license-file: LICENSE +author: Alexander Thiemann +maintainer: mail@athiemann.net +copyright: 2017 Alexander Thiemann +category: Web +build-type: Simple +extra-source-files: README.md +cabal-version: >=1.10 + +library + hs-source-dirs: src + exposed-modules: Data.FilePlow + , Data.FilePlow.Ordered + build-depends: base >= 4.7 && < 5 + , bytestring + , vector + , binary-search + default-language: Haskell2010 + ghc-options: -Wall + +test-suite fileplow-test + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: Spec.hs + build-depends: base + , fileplow + , hspec + , temporary + , bytestring + , QuickCheck >= 2.9 + , mtl + ghc-options: -threaded -rtsopts -with-rtsopts=-N + default-language: Haskell2010 + +source-repository head + type: git + location: https://github.com/agrafix/fileplow diff --git a/src/Data/FilePlow.hs b/src/Data/FilePlow.hs new file mode 100644 --- /dev/null +++ b/src/Data/FilePlow.hs @@ -0,0 +1,183 @@ +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE BangPatterns #-} +module Data.FilePlow + ( PlowHandle(..), Handle, SeekMode(..) + , seekUntil, seekUntilRev + , MultiHandle, withMultiHandle + ) +where + +import Control.Monad +import Data.IORef +import GHC.IO.Exception +import GHC.IO.Handle +import System.IO +import qualified Data.ByteString as BS +import qualified Data.Vector as V + +class PlowHandle hdl where + pSeek :: hdl -> SeekMode -> Integer -> IO () + pTell :: hdl -> IO Integer + pGetChar :: hdl -> IO Char + pGetLine :: hdl -> IO BS.ByteString + pIsEOF :: hdl -> IO Bool + pFileSize :: hdl -> IO Integer + +instance PlowHandle Handle where + pSeek = hSeek + pTell = hTell + pGetChar = hGetChar + pGetLine = BS.hGetLine + pIsEOF = hIsEOF + pFileSize = hFileSize + +data MultiHandle h + = MultiHandle + { mh_handles :: !(V.Vector h) + , mh_currentIndex :: !(IORef Int) + } + +-- | Seek until a certain charater is reached in reverse +seekUntilRev :: PlowHandle hdl => hdl -> (Char -> Bool) -> IO Bool +seekUntilRev hdl pp = + do let getLoop = + do p <- pTell hdl + if p == 0 + then pure False + else do x <- pGetChar hdl + if pp x + then pure True + else do pSeek hdl RelativeSeek (-2) + getLoop + getLoop + +-- | Seek until a certain charater is reached +seekUntil :: PlowHandle hdl => hdl -> (Char -> Bool) -> IO Bool +seekUntil hdl pp = + do let getLoop = + do eof <- pIsEOF hdl + if eof + then pure False + else do x <- pGetChar hdl + if pp x then pure True else getLoop + getLoop + +instance (PlowHandle h) => PlowHandle (MultiHandle h) where + pFileSize mh = + V.foldM' (\total hdl -> (+ total) <$> pFileSize hdl) 0 (mh_handles mh) + {-# INLINE pFileSize #-} + + pIsEOF mh = + do idx <- readIORef (mh_currentIndex mh) + let ct = V.length (mh_handles mh) + if idx >= ct - 1 + then pIsEOF $ mh_handles mh V.! (ct - 1) + else pure False + {-# INLINE pIsEOF #-} + + pGetChar mh = + do h <- getCurrentHandle "pGetChar" mh + pGetChar h + {-# INLINE pGetChar #-} + + pGetLine mh = + do h <- getCurrentHandle "pGetLine" mh + pGetLine h + {-# INLINE pGetLine #-} + + pSeek mh sm val = + case sm of + AbsoluteSeek -> absoluteSeek mh val + RelativeSeek -> relativeSeek mh val + SeekFromEnd -> seekFromEnd mh val + {-# INLINE pSeek #-} + + pTell mh = + do idx <- readIORef (mh_currentIndex mh) + let h = mh_handles mh V.! idx + localPos <- pTell h + foldM + (\total i -> (+ total) <$> pFileSize (mh_handles mh V.! i)) + localPos [0 .. (idx - 1)] + {-# INLINE pTell #-} + +seekFromEnd :: (PlowHandle h) => MultiHandle h -> Integer -> IO () +seekFromEnd mh tval = + let ct = V.length (mh_handles mh) + seekLoop !tgtVal !idx + | idx >= ct || idx < 0 = + ioException (IOError Nothing EOF "seekFromEnd" "No more file handles" Nothing Nothing) + | otherwise = + do let h = mh_handles mh V.! idx + size <- pFileSize h + if abs tgtVal >= size + then seekLoop ((-1) * (abs tgtVal - size)) (idx - 1) + else do writeIORef (mh_currentIndex mh) idx + pSeek h SeekFromEnd tgtVal + in seekLoop tval (ct - 1) + +relativeSeek :: (PlowHandle h) => MultiHandle h -> Integer -> IO () +relativeSeek mh tval = + do myIdx <- readIORef (mh_currentIndex mh) + let localH = mh_handles mh V.! myIdx + localPos <- pTell localH + pSeek localH AbsoluteSeek 0 + let seekVal = tval + localPos + seekHelp mh seekVal myIdx + +absoluteSeek :: PlowHandle h => MultiHandle h -> Integer -> IO () +absoluteSeek mh tval = + seekHelp mh tval 0 + +seekHelp :: PlowHandle h => MultiHandle h -> Integer -> Int -> IO () +seekHelp mh t i = + let ct = V.length (mh_handles mh) + seekLoop !tgtVal !idx + | idx >= ct = + let msg = "No more file handles, desired pos: " ++ show t + in ioException (IOError Nothing EOF "seekHelp" msg Nothing Nothing) + | tgtVal < 0 && idx > 0 = + do let h = mh_handles mh V.! (idx - 1) + size <- pFileSize h + pSeek h AbsoluteSeek 0 + seekLoop (tgtVal + size) (idx - 1) + | otherwise = + do let h = mh_handles mh V.! idx + size <- pFileSize h + if tgtVal >= size + then seekLoop (tgtVal - size) (idx + 1) + else do writeIORef (mh_currentIndex mh) idx + pSeek h AbsoluteSeek tgtVal + in seekLoop t i + +getCurrentHandle :: PlowHandle h => String -> MultiHandle h -> IO h +getCurrentHandle helpTxt mh = + do startIdx <- readIORef (mh_currentIndex mh) + let ct = V.length (mh_handles mh) + gotoUsefulHdl !i !nh + | i >= ct = + ioException (IOError Nothing EOF ("currentHandle/" ++ helpTxt) "No more file handles" Nothing Nothing) + | otherwise = + do let h = mh_handles mh V.! i + when nh $ + pSeek h AbsoluteSeek 0 + eof <- pIsEOF h + if eof + then gotoUsefulHdl (i + 1) True + else do writeIORef (mh_currentIndex mh) i + pure h + gotoUsefulHdl startIdx False + + +withMultiHandle :: [FilePath] -> (MultiHandle Handle -> IO a) -> IO a +withMultiHandle files go = + loop [] files + where + loop accum xs = + case xs of + [] -> + do r <- newIORef 0 + go (MultiHandle (V.fromList (reverse accum)) r) + (fp : more) -> + withFile fp ReadMode $ \hdl -> + loop (hdl : accum) more diff --git a/src/Data/FilePlow/Ordered.hs b/src/Data/FilePlow/Ordered.hs new file mode 100644 --- /dev/null +++ b/src/Data/FilePlow/Ordered.hs @@ -0,0 +1,46 @@ +module Data.FilePlow.Ordered + ( SeekTarget(..), seekTo ) +where + +import Data.FilePlow + +import Numeric.Search + +data SeekTarget b + = StLargerThan !b + -- ^ put the handle to the first entry that is larger than the provided value. + -- Useful for files with entries in ASCending order + | StSmallerThan !b + -- ^ put the handle to the first entry that is smaller than the provided value. + -- Useful for files with entries in DESCending order + deriving (Show, Eq) + +-- ^ place the file handling to a specific 'SeekTarget' provided a function that +-- extracts a value/entry at a specific position +seekTo :: + (Ord b, Eq b, PlowHandle hdl) + => hdl + -> SeekTarget b + -> (hdl -> IO b) + -> IO Bool +seekTo hdl st extract = + do s <- pFileSize hdl + ranges <- searchM (fromTo 0 (s - 1)) divForever $ \pos -> + do case st of + StLargerThan _ -> pSeek hdl AbsoluteSeek pos + StSmallerThan _ -> pSeek hdl AbsoluteSeek (s - 1 - pos) + val <- extract hdl + pure $ + case st of + StLargerThan x -> (val >= x) + StSmallerThan x -> (val <= x) + let tgtVal = + case st of + StSmallerThan _ -> largest True ranges + StLargerThan _ -> smallest True ranges + case tgtVal of + Just pos -> + case st of + StLargerThan _ -> pSeek hdl AbsoluteSeek pos >> pure True + StSmallerThan _ -> pSeek hdl AbsoluteSeek (s - 1 - pos) >> pure True + Nothing -> pure False diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,157 @@ +module Main where + +import Data.FilePlow +import Data.FilePlow.Ordered + +import Control.Monad +import Control.Monad.Trans +import Data.List +import Data.Monoid +import System.IO (hClose) +import System.IO.Temp +import Test.Hspec +import Test.QuickCheck +import Test.QuickCheck.Monadic +import qualified Data.ByteString as BS +import qualified Data.ByteString.Builder as BSB +import qualified Data.ByteString.Char8 as BSC +import qualified Data.ByteString.Lazy as BSL + +makeRangeBS :: ([Int] -> [Int]) -> Int -> Int -> BS.ByteString +makeRangeBS f lb ub = + BSL.toStrict . BSB.toLazyByteString $ + foldl' (\bldr i -> bldr <> BSB.int64Dec (fromIntegral i) <> BSB.char8 '\n') mempty $ f [lb .. ub] + +withDummyFileNum :: ([Int] -> [Int]) -> (Int, Int) -> (FilePath -> IO a) -> IO a +withDummyFileNum f (lb, ub) go = + withSystemTempFile "dummyFilePlowTest" $ \fp hdl -> + do BS.hPut hdl (makeRangeBS f lb ub) + hClose hdl + go fp + +withDummyFile :: Char -> Int -> (FilePath -> IO a) -> IO a +withDummyFile c size go = + withSystemTempFile "dummyFilePlowTest" $ \fp hdl -> + do BS.hPut hdl (BSC.replicate size c) + hClose hdl + go fp + +readRest :: MultiHandle Handle -> IO BS.ByteString +readRest hdl = + do let getLoop builder = + do eof <- pIsEOF hdl + if eof + then pure $ BSL.toStrict (BSB.toLazyByteString builder) + else do x <- pGetChar hdl + getLoop (builder <> BSB.char8 x) + getLoop mempty + +main :: IO () +main = + hspec $ + do describe "multi handle" $ + do it "simple get char works" $ + withDummyFile 'a' 100 $ \f1 -> + withDummyFile 'b' 100 $ \f2 -> + withMultiHandle [f1, f2] $ \hdl -> + do bs <- readRest hdl + bs `shouldBe` (BSC.replicate 100 'a' <> BSC.replicate 100 'b') + it "two full reads with going to the begining work" $ + withDummyFile 'a' 100 $ \f1 -> + withDummyFile 'b' 100 $ \f2 -> + withMultiHandle [f1, f2] $ \hdl -> + do _ <- readRest hdl + pSeek hdl AbsoluteSeek 0 + bs <- readRest hdl + bs `shouldBe` (BSC.replicate 100 'a' <> BSC.replicate 100 'b') + it "seek tell works" $ + forAll (choose (0, 199)) $ \p -> + monadicIO $ + do r <- + liftIO $ + withDummyFile 'a' 100 $ \f1 -> + withDummyFile 'b' 100 $ \f2 -> + withMultiHandle [f1, f2] $ \hdl -> + do pSeek hdl AbsoluteSeek p + x <- pTell hdl + pure x + assert (r == p) + it "seek until works" $ + withDummyFile 'a' 100 $ \f1 -> + withDummyFile 'b' 100 $ \f2 -> + withMultiHandle [f1, f2] $ \hdl -> + do ok <- seekUntil hdl (== 'b') + ok `shouldBe` True + bs <- readRest hdl + bs `shouldBe` BSC.replicate 99 'b' + it "seek untilRev works" $ + withDummyFile 'a' 100 $ \f1 -> + withDummyFile 'b' 100 $ \f2 -> + withMultiHandle [f1, f2] $ \hdl -> + do pSeek hdl AbsoluteSeek 199 + ok <- seekUntilRev hdl (== 'a') + ok `shouldBe` True + bs <- readRest hdl + bs `shouldBe` BSC.replicate 100 'b' + it "seeked reads work" $ + withDummyFile 'a' 100 $ \f1 -> + withDummyFile 'b' 100 $ \f2 -> + withMultiHandle [f1, f2] $ \hdl -> + do pSeek hdl AbsoluteSeek 100 + bs <- readRest hdl + bs `shouldBe` BSC.replicate 100 'b' + + pSeek hdl AbsoluteSeek 50 + bs2 <- readRest hdl + bs2 `shouldBe` (BSC.replicate 50 'a' <> BSC.replicate 100 'b') + + pSeek hdl SeekFromEnd (-50) + bs3 <- readRest hdl + bs3 `shouldBe` BSC.replicate 50 'b' + + pSeek hdl SeekFromEnd (-150) + bs4 <- readRest hdl + bs4 `shouldBe` (BSC.replicate 50 'a' <> BSC.replicate 100 'b') + + pSeek hdl AbsoluteSeek 0 + pSeek hdl RelativeSeek 100 + bs5 <- readRest hdl + bs5 `shouldBe` BSC.replicate 100 'b' + + pSeek hdl AbsoluteSeek 50 + pSeek hdl RelativeSeek 25 + bs6 <- readRest hdl + bs6 `shouldBe` (BSC.replicate 25 'a' <> BSC.replicate 100 'b') + describe "ordered" $ + do it "find a lower bound" $ + withDummyFileNum id (100, 199) $ \f1 -> + withDummyFileNum id (200, 299) $ \f2 -> + withMultiHandle [f1, f2] $ \hdl -> + do ok <- + seekTo hdl (StLargerThan (150 :: Int)) $ \h -> + do void $ seekUntilRev hdl (== '\n') + eof <- pIsEOF hdl + if eof + then pure 999 + else do ln <- pGetLine h + pure $ read (BSC.unpack ln) + ok `shouldBe` True + void $ seekUntilRev hdl (== '\n') + bs <- readRest hdl + bs `shouldBe` (makeRangeBS id 150 299) + it "finds upper bound" $ + withDummyFileNum reverse (200, 299) $ \f1 -> + withDummyFileNum reverse (100, 199) $ \f2 -> + withMultiHandle [f1, f2] $ \hdl -> + do ok <- + seekTo hdl (StSmallerThan (150 :: Int)) $ \h -> + do void $ seekUntilRev hdl (== '\n') + eof <- pIsEOF hdl + if eof + then pure 0 + else do ln <- pGetLine h + pure $ read (BSC.unpack ln) + ok `shouldBe` True + void $ seekUntilRev hdl (== '\n') + bs <- readRest hdl + bs `shouldBe` (makeRangeBS reverse 100 150)