diff --git a/ChangeLog.md b/ChangeLog.md --- a/ChangeLog.md +++ b/ChangeLog.md @@ -1,3 +1,7 @@ +# 0.1.0.2 -> 0.1.0.3 +- Changed bytestring-mmap to mmap +- Add LRU caching to improve performance (10000 entries are cached) + # 0.1.0.1 -> 0.1.0.2 - GHC-7.10 compatibiliy diff --git a/Data/GeoIP2.hs b/Data/GeoIP2.hs --- a/Data/GeoIP2.hs +++ b/Data/GeoIP2.hs @@ -1,6 +1,6 @@ -{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE MultiWayIf #-} +{-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE MultiWayIf #-} module Data.GeoIP2 ( -- * Library description @@ -23,32 +23,35 @@ , GeoResult(..) ) where -import Control.Monad (when, unless) -import System.IO.Posix.MMap -import qualified Data.ByteString as BS -import qualified Data.ByteString.Lazy as BL -import Data.Int -import Data.Binary -import qualified Data.Text as T -import Data.IP (IP(..), ipv4ToIPv6) -import Control.Applicative ((<$>), (<*>)) -import qualified Data.Map as Map -import Data.Maybe (mapMaybe, fromMaybe) +import Control.Applicative ((<$>), (<*>)) +import Control.Monad (unless, when) +import Data.Binary +import qualified Data.ByteString as BS +import qualified Data.ByteString.Lazy as BL +import qualified Data.Cache.LRU.IO as LRU +import Data.Int +import Data.IP (IP (..), ipv4ToIPv6) +import qualified Data.Map as Map +import Data.Maybe (fromMaybe, mapMaybe) +import qualified Data.Text as T +import System.IO.MMap +import System.IO.Unsafe (unsafePerformIO) -import Data.GeoIP2.Fields -import Data.GeoIP2.SearchTree +import Data.GeoIP2.Fields +import Data.GeoIP2.SearchTree -- | Address type stored in database data GeoIP = GeoIPv6 | GeoIPv4 deriving (Eq, Show) -- | Handle for search operations data GeoDB = GeoDB { - geoMem :: BL.ByteString - , geoDbType :: T.Text -- ^ String that indicates the structure of each data record associated with an IP address - , geoDbLanguages :: [T.Text] -- ^ Languages supported in database - , geoDbNodeCount :: Int64 - , geoDbRecordSize :: Int - , geoDbAddrType :: GeoIP -- ^ Type of address (IPv4/IPv6) stored in a database + geoMem :: BS.ByteString + , geoLru :: LRU.AtomicLRU Int GeoField + , geoDbType :: T.Text -- ^ String that indicates the structure of each data record associated with an IP address + , geoDbLanguages :: [T.Text] -- ^ Languages supported in database + , geoDbNodeCount :: Int64 + , geoDbRecordSize :: Int + , geoDbAddrType :: GeoIP -- ^ Type of address (IPv4/IPv6) stored in a database , geoDbDescription :: Maybe T.Text -- ^ Description of a database in english } @@ -63,12 +66,12 @@ -- | Open database, mmap it into memory, parse header and return a handle for search operations openGeoDB :: FilePath -> IO GeoDB openGeoDB geoFile = do - bsmem <- unsafeMMapFile geoFile - let (DataMap hdr) = decode (BL.fromChunks [getHeaderBytes bsmem]) - mem = BL.fromChunks [bsmem] + bsmem <- mmapFileByteString geoFile Nothing + let (DataMap hdr) = decode (BL.fromStrict $ getHeaderBytes bsmem) when (hdr .: "binary_format_major_version" /= (2 :: Int)) $ error "Unsupported database version, only v2 supported." unless (hdr .: "record_size" `elem` [24, 28, 32 :: Int]) $ error "Record size not supported." - return $ GeoDB mem (hdr .: "database_type") + lru <- LRU.newAtomicLRU (Just 10000) + return $ GeoDB bsmem lru (hdr .: "database_type") (fromMaybe [] $ hdr .:? "languages") (hdr .: "node_count") (hdr .: "record_size") (if (hdr .: "ip_version") == (4 :: Int) then GeoIPv4 else GeoIPv6) @@ -79,35 +82,46 @@ rawGeoData :: Monad m => GeoDB -> IP -> m GeoField rawGeoData geodb addr = do bits <- coerceAddr - offset <- getDataOffset (geoMem geodb, geoDbNodeCount geodb, geoDbRecordSize geodb) bits - let basedata = dataAt offset - return $ resolvePointers basedata + offset <- fromIntegral <$> getDataOffset (geoMem geodb, geoDbNodeCount geodb, geoDbRecordSize geodb) bits + basedata <- dataAt offset + resolvePointers basedata where - dataSectionStart = fromIntegral (geoDbRecordSize geodb `div` 4) * geoDbNodeCount geodb + 16 - dataAt offset = decode (BL.drop (offset + dataSectionStart) (geoMem geodb)) + dataSectionStart = (geoDbRecordSize geodb `div` 4) * fromIntegral (geoDbNodeCount geodb) + 16 + -- Add caching + dataAt offset = case lruget offset of + Right res -> return res + Left err -> fail err + lruget offset = unsafePerformIO $ do + cached <- LRU.lookup offset (geoLru geodb) + case cached of + Just result -> return (Right result) + Nothing -> do + let decres = decodeOrFail (BL.fromStrict $ BS.drop (offset + dataSectionStart) (geoMem geodb)) + case decres of + Left (_,_,err) -> return (Left err) + Right (_,_,result) -> do + LRU.insert offset result (geoLru geodb) + return (Right result) coerceAddr | (IPv4 _) <- addr, GeoIPv4 <- geoDbAddrType geodb = return $ ipToBits addr | (IPv6 _) <- addr, GeoIPv6 <- geoDbAddrType geodb = return $ ipToBits addr | (IPv4 addrv4) <- addr, GeoIPv6 <- geoDbAddrType geodb = return $ ipToBits $ IPv6 (ipv4ToIPv6 addrv4) | otherwise = fail "Cannot search IPv6 address in IPv4 database" - resolvePointers (DataPointer ptr) = resolvePointers $ dataAt ptr - resolvePointers (DataMap obj) = DataMap $ Map.fromList $ map resolveTuple (Map.toList obj) - resolvePointers (DataArray arr) = DataArray $ map resolvePointers arr - resolvePointers x = x - resolveTuple (a,b) = - let r1 = resolvePointers a - r2 = resolvePointers b - in (r1, r2) + resolvePointers (DataPointer ptr) = dataAt (fromIntegral ptr) >>= resolvePointers -- TODO - limit recursion? + resolvePointers (DataMap obj) = DataMap . Map.fromList <$> mapM resolveTuple (Map.toList obj) + resolvePointers (DataArray arr) = DataArray <$> mapM resolvePointers arr + resolvePointers x = return x + resolveTuple (a,b) = (,) <$> resolvePointers a <*> resolvePointers b -- | Result of a search query data GeoResult = GeoResult { - geoContinent :: Maybe T.Text + geoContinent :: Maybe T.Text , geoContinentCode :: Maybe T.Text - , geoCountryISO :: Maybe T.Text - , geoCountry :: Maybe T.Text - , geoLocation :: Maybe (Double, Double) - , geoCity :: Maybe T.Text - , geoSubdivisions :: [(T.Text, T.Text)] + , geoCountryISO :: Maybe T.Text + , geoCountry :: Maybe T.Text + , geoLocation :: Maybe (Double, Double) + , geoCity :: Maybe T.Text + , geoSubdivisions :: [(T.Text, T.Text)] } deriving (Show, Eq) -- | Search GeoIP database, monadic version (e.g. use with Maybe or Either) diff --git a/Data/GeoIP2/Fields.hs b/Data/GeoIP2/Fields.hs --- a/Data/GeoIP2/Fields.hs +++ b/Data/GeoIP2/Fields.hs @@ -1,21 +1,22 @@ {-# OPTIONS_HADDOCK hide #-} -{-# LANGUAGE MultiWayIf #-} {-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiWayIf #-} module Data.GeoIP2.Fields where import Control.Applicative ((<$>)) -import Control.Monad (replicateM, replicateM_) +import Control.Monad (replicateM) import Data.Binary import Data.Binary.Get import Data.Bits (shift, (.&.)) +import qualified Data.ByteString as BS import Data.Int import qualified Data.Map as Map import Data.Maybe (fromMaybe) import Data.ReinterpretCast (wordToDouble) import qualified Data.Text as T import Data.Text.Encoding (decodeUtf8) -import Data.Word () +import Data.Word () data GeoField = DataPointer Int64 @@ -76,8 +77,8 @@ -- | Parse number of given length parseNumber :: Num a => Int64 -> Get a parseNumber fsize = do - bytes <- map fromIntegral <$> replicateM (fromIntegral fsize) getWord8 - return $ foldl (\acc new -> new + 256 * acc) 0 bytes + bytes <- getByteString (fromIntegral fsize) + return $ BS.foldl' (\acc new -> fromIntegral new + 256 * acc) 0 bytes instance Binary GeoField where put = undefined @@ -118,5 +119,5 @@ 11 -> DataArray <$> replicateM (fromIntegral fsize) get 14 -> return $ DataBool (fsize == 0) _ -> do - replicateM_ (fromIntegral fsize) getWord8 + _ <- getByteString (fromIntegral fsize) return $ DataUnknown ftype fsize diff --git a/Data/GeoIP2/SearchTree.hs b/Data/GeoIP2/SearchTree.hs --- a/Data/GeoIP2/SearchTree.hs +++ b/Data/GeoIP2/SearchTree.hs @@ -5,7 +5,6 @@ import Data.Binary import Data.Bits (shift, testBit, (.&.), (.|.)) import qualified Data.ByteString as BS -import qualified Data.ByteString.Lazy as BL import Data.Int import Data.IP (IP (..), fromIPv4, fromIPv6b) @@ -19,14 +18,12 @@ ipToBits (IPv6 addr) = concatMap byteToBits (fromIPv6b addr) -- | Read node (2 records) given the index of a node -readNode :: BL.ByteString -> Int -> Int64 -> (Int64, Int64) +readNode :: BS.ByteString -> Int -> Int64 -> (Int64, Int64) readNode mem recordbits index = let bytecount = fromIntegral $ recordbits `div` 4 - bytes = BL.take (fromIntegral bytecount) $ BL.drop (fromIntegral $ index * bytecount) mem - numbers = concatMap BS.unpack (BL.toChunks bytes) :: [Word8] - makenum = foldl (\acc new -> fromIntegral new + 256 * acc) 0 :: [Word8] -> Word64 - num = makenum numbers + bytes = BS.take (fromIntegral bytecount) $ BS.drop (fromIntegral $ index * bytecount) mem + num = BS.foldl' (\acc new -> fromIntegral new + 256 * acc) 0 bytes :: Word64 -- 28 bits has a strange record format left28 = num `shift` (-32) .|. (num .&. 0xf0000000) in case recordbits of @@ -34,7 +31,7 @@ _ -> (fromIntegral (num `shift` negate recordbits), fromIntegral (num .&. ((1 `shift` recordbits) - 1))) -- | Get offset in the Data Section -getDataOffset :: Monad m => (BL.ByteString, Int64, Int) -> [Bool] -> m Int64 +getDataOffset :: Monad m => (BS.ByteString, Int64, Int) -> [Bool] -> m Int64 getDataOffset (mem, nodeCount, recordSize) startbits = getnode startbits 0 where diff --git a/geoip2.cabal b/geoip2.cabal --- a/geoip2.cabal +++ b/geoip2.cabal @@ -1,5 +1,5 @@ name: geoip2 -version: 0.1.0.2 +version: 0.1.0.3 synopsis: Pure haskell interface to MaxMind GeoIP database description: GeoIP2 is a haskell binding to the MaxMind GeoIP2 database. @@ -24,12 +24,13 @@ exposed-modules: Data.GeoIP2 other-modules: Data.GeoIP2.Fields, Data.GeoIP2.SearchTree build-depends: base >=4.7 && <4.9 - , bytestring-mmap + , mmap , bytestring , binary , text , containers , iproute(>=1.4.0) , reinterpret-cast + , lrucache default-language: Haskell2010 ghc-options: -Wall