diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,29 @@ +Copyright (c) 2014, IIJ Innovation Institute Inc. +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 the copyright holders nor the names of its + 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/Network/DNS/Cache.hs b/Network/DNS/Cache.hs new file mode 100644 --- /dev/null +++ b/Network/DNS/Cache.hs @@ -0,0 +1,256 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | DNS cache to resolve domains concurrently. + +module Network.DNS.Cache ( + DNSCacheConf(..) + , DNSCache + , withDNSCache + -- * Looking up + , lookup + , lookupCache + -- * Resolving + , Result(..) + , resolve + , resolveCache + -- * Waiting + , wait + ) where + +import Control.Applicative ((<$>)) +import Control.Concurrent (threadDelay, forkIO) +import Control.Concurrent.Async (async, waitAnyCancel) +import Control.Exception (bracket) +import Control.Monad (forever, void) +import qualified Data.ByteString.Char8 as BS +import Data.IP (toHostAddress) +import Data.Time (getCurrentTime, addUTCTime, NominalDiffTime) +import Network.DNS hiding (lookup) +import Network.DNS.Cache.Cache +import qualified Network.DNS.Cache.Sync as S +import Network.DNS.Cache.Types +import Network.DNS.Cache.Utils +import Network.DNS.Cache.Value +import Prelude hiding (lookup) + +---------------------------------------------------------------- + +-- | Configuration for DNS cache. +data DNSCacheConf = DNSCacheConf { + -- | A list of resolvers (cache DNS servers). + -- A domain is resolved by the resolvers concurrently. + -- The first reply is used regardless of success/failure at this moment + resolvConfs :: [ResolvConf] + -- | Capability of how many domains can be resolved concurrently + , maxConcurrency :: Int + -- | The minimum bound of cache duration for success replies in seconds. + , minTTL :: NominalDiffTime + -- | The maximum bound of cache duration for success replies in seconds. + , maxTTL :: NominalDiffTime + -- | The cache duration for failure replies in seconds. + , negativeTTL :: NominalDiffTime + } + +-- | An abstract data for DNS cache. +-- Cached domains are expired every 10 seconds according to their TTL. +data DNSCache = DNSCache { + cacheSeeds :: [ResolvSeed] + , cacheNofServers :: !Int + , cacheRef :: CacheRef + , cacheActiveRef :: S.ActiveRef + , cacheConcVar :: S.ConcVar + , cacheConcLimit :: Int + , cacheMinTTL :: NominalDiffTime + , cacheMaxTTL :: NominalDiffTime + , cacheNegTTL :: NominalDiffTime + } + +---------------------------------------------------------------- + +-- | A basic function to create DNS cache. +-- Domains should be resolved in the function of the second argument. +withDNSCache :: DNSCacheConf -> (DNSCache -> IO a) -> IO a +withDNSCache conf func = do + seeds <- mapM makeResolvSeed (resolvConfs conf) + let n = length seeds + cacheref <- newCacheRef + activeref <- S.newActiveRef + lvar <- S.newConcVar + let cache = DNSCache seeds n cacheref activeref lvar maxcon minttl maxttl negttl + void . forkIO $ prune cacheref + func cache + where + maxcon = maxConcurrency conf + minttl = minTTL conf + maxttl = maxTTL conf + negttl = negativeTTL conf + +---------------------------------------------------------------- + +lookupPSQ :: DNSCache -> Domain -> IO (Key, Maybe (Prio, Entry)) +lookupPSQ cache dom = do + !mx <- lookupCacheRef key cacheref + return (key,mx) + where + cacheref = cacheRef cache + !key = newKey dom + +---------------------------------------------------------------- + +-- | Lookup 'Domain' only in the cache. +lookupCache :: DNSCache -> Domain -> IO (Maybe HostAddress) +lookupCache cache dom = do + mx <- resolveCache cache dom + case mx of + Nothing -> return Nothing + Just ev -> return (fromEither ev) + +---------------------------------------------------------------- + +-- | Lookup 'Domain' in the cache. +-- If not exist, queries are sent to DNS servers and +-- resolved IP addresses are cached. +lookup :: DNSCache -> Domain -> IO (Maybe HostAddress) +lookup cache dom = fromEither <$> resolve cache dom + +---------------------------------------------------------------- + +-- | Lookup 'Domain' only in the cache. +resolveCache :: DNSCache -> Domain -> IO (Maybe (Either DNSError Result)) +resolveCache _ dom + | isIPAddr dom = Just . Right . Numeric <$> return (tov4 dom) + where + tov4 = read . BS.unpack +resolveCache cache dom = do + (_, mx) <- lookupPSQ cache dom + case mx of + Nothing -> return Nothing + Just (_, Right v) -> Just . Right . Hit <$> rotate v + Just (_, Left e) -> Just . Left <$> return e + +-- | Lookup 'Domain' in the cache. +-- If not exist, queries are sent to DNS servers and +-- resolved IP addresses are cached. +resolve :: DNSCache -> Domain -> IO (Either DNSError Result) +resolve _ dom + | isIPAddr dom = return $ Right $ Numeric $ tov4 dom + where + tov4 = read . BS.unpack +resolve cache dom = do + (key,mx) <- lookupPSQ cache dom + case mx of + Just (_,ev) -> case ev of + Left e -> Left <$> return e + Right v -> Right . Hit <$> rotate v + Nothing -> do + -- If this domain is being resolved by another thread + -- let's wait. + ma <- S.lookupActiveRef key activeref + case ma of + Just avar -> S.listen avar + Nothing -> do + avar <- S.newActiveVar + S.insertActiveRef key avar activeref + x <- sendQuery cache dom + !res <- case x of + Left err -> insertNegative cache key err + Right [] -> insertNegative cache key UnexpectedRDATA + Right addrs -> insertPositive cache key addrs + S.deleteActiveRef key activeref + S.tell avar res + return res + where + activeref = cacheActiveRef cache + +insertPositive :: DNSCache -> Key -> [(HostAddress, TTL)] + -> IO (Either DNSError Result) +insertPositive _ _ [] = error "insertPositive" +insertPositive cache key addrs@((addr,ttl):_) = do + !ent <- positiveEntry $ map fst addrs + !tim <- addUTCTime lifeTime <$> getCurrentTime + insertCacheRef key tim ent cacheref + return $! Right $ Resolved addr + where + minttl = cacheMinTTL cache + maxttl = cacheMaxTTL cache + !lifeTime = minttl `max` (maxttl `min` fromIntegral ttl) + cacheref = cacheRef cache + +insertNegative :: DNSCache -> Key -> DNSError -> IO (Either DNSError Result) +insertNegative cache key err = do + !tim <- addUTCTime lifeTime <$> getCurrentTime + insertCacheRef key tim (Left err) cacheref + return $ Left err + where + lifeTime = cacheNegTTL cache + cacheref = cacheRef cache + +---------------------------------------------------------------- + +sendQuery :: DNSCache -> Domain -> IO (Either DNSError [(HostAddress,TTL)]) +sendQuery cache dom = bracket setup teardown body + where + setup = waitIncrease cache + teardown _ = decrease cache + body _ = concResolv cache dom + +waitIncrease :: DNSCache -> IO () +waitIncrease cache = S.waitIncrease lvar lim + where + lvar = cacheConcVar cache + lim = cacheConcLimit cache + +decrease :: DNSCache -> IO () +decrease cache = S.decrease lvar + where + lvar = cacheConcVar cache + +concResolv :: DNSCache -> Domain -> IO (Either DNSError [(HostAddress,TTL)]) +concResolv cache dom = withResolvers seeds $ \resolvers -> do + eans <- resolv n resolvers dom + return $ case eans of + Left err -> Left err + Right ans -> fromDNSFormat ans getHostAddressandTTL + where + n = cacheNofServers cache + seeds = cacheSeeds cache + isA r = rrtype r == A + unTag (RD_A ip) = ip + unTag _ = error "unTag" + toAddr = toHostAddress . unTag . rdata + hostAddressandTTL r = (toAddr r, rrttl r) + getHostAddressandTTL = map hostAddressandTTL . filter isA . answer + +resolv :: Int -> [Resolver] -> Domain -> IO (Either DNSError DNSFormat) +resolv 1 resolvers dom = lookupRaw (head resolvers) dom A +resolv _ resolvers dom = do + asyncs <- mapM async actions + snd <$> waitAnyCancel asyncs + where + actions = map (\res -> lookupRaw res dom A) resolvers + +---------------------------------------------------------------- + +-- | Wait until the predicate in the second argument is satisfied. +-- The predicate are given the number of the current resolving domains. +-- +-- For instance, if you ensure that no resolvings are going on: +-- +-- > wait cache (== 0) +-- +-- If you want to ensure that capability of concurrent resolving is not full: +-- +-- > wait cache (< maxCon) +-- +-- where 'maxCon' represents 'maxConcurrency' in 'DNSCacheConf'. +wait :: DNSCache -> (Int -> Bool) -> IO () +wait cache cond = S.wait lvar cond + where + lvar = cacheConcVar cache + +prune :: CacheRef -> IO () +prune cacheref = forever $ do + threadDelay 10000000 + tim <- getCurrentTime + pruneCacheRef tim cacheref diff --git a/Network/DNS/Cache/Cache.hs b/Network/DNS/Cache/Cache.hs new file mode 100644 --- /dev/null +++ b/Network/DNS/Cache/Cache.hs @@ -0,0 +1,41 @@ +{-# LANGUAGE BangPatterns #-} + +module Network.DNS.Cache.Cache ( + CacheRef + , newCacheRef + , lookupCacheRef + , insertCacheRef + , pruneCacheRef + , newKey + ) where + +import Control.Applicative ((<$>)) +import Data.ByteString (ByteString) +import qualified Data.ByteString.Short as B +import Data.Hashable (hash) +import Data.IORef (newIORef, readIORef, atomicModifyIORef', IORef) +import Network.DNS.Cache.PSQ (PSQ) +import qualified Network.DNS.Cache.PSQ as PSQ +import Network.DNS.Cache.Types + +newtype CacheRef = CacheRef (IORef (PSQ Entry)) + +newCacheRef :: IO CacheRef +newCacheRef = CacheRef <$> newIORef PSQ.empty + +lookupCacheRef :: Key -> CacheRef -> IO (Maybe (Prio, Entry)) +lookupCacheRef key (CacheRef ref) = PSQ.lookup key <$> readIORef ref + +insertCacheRef :: Key -> Prio -> Entry -> CacheRef -> IO () +insertCacheRef key tim ent (CacheRef ref) = + atomicModifyIORef' ref $ \q -> (PSQ.insert key tim ent q, ()) + +pruneCacheRef :: Prio -> CacheRef -> IO () +pruneCacheRef tim (CacheRef ref) = + atomicModifyIORef' ref $ \p -> (snd (PSQ.atMost tim p), ()) + +newKey :: ByteString -> Key +newKey dom = Key h k + where + !k = B.toShort dom + !h = hash dom diff --git a/Network/DNS/Cache/PSQ.hs b/Network/DNS/Cache/PSQ.hs new file mode 100644 --- /dev/null +++ b/Network/DNS/Cache/PSQ.hs @@ -0,0 +1,480 @@ +{-# LANGUAGE BangPatterns #-} + +-- From ghc-7.6.3/libraries/base/GHC/Event/PSQ.hs + +-- Copyright (c) 2008, Ralf Hinze +-- 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. +-- +-- * The names of the contributors may not 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. + +-- | A /priority search queue/ (henceforth /queue/) efficiently +-- supports the operations of both a search tree and a priority queue. +-- An 'Elem'ent is a product of a key, a priority, and a +-- value. Elements can be inserted, deleted, modified and queried in +-- logarithmic time, and the element with the least priority can be +-- retrieved in constant time. A queue can be built from a list of +-- elements, sorted by keys, in linear time. +-- +-- This implementation is due to Ralf Hinze with some modifications by +-- Scott Dillard and Johan Tibell. +-- +-- * Hinze, R., /A Simple Implementation Technique for Priority Search +-- Queues/, ICFP 2001, pp. 110-121 +-- +-- +module Network.DNS.Cache.PSQ + ( + -- * Binding Type + Elem(..) + , Key + , Prio + + -- * Priority Search Queue Type + , PSQ + + -- * Query + , size + , null + , lookup + + -- * Construction + , empty + , singleton + + -- * Insertion + , insert + + -- * Delete/Update + , delete + , adjust + + -- * Conversion + , toList + , toAscList + , toDescList + , fromList + + -- * Min + , findMin + , deleteMin + , minView + , atMost + ) where + +import Network.DNS.Cache.Types +import Prelude hiding (lookup, null) + +-- | @E k p@ binds the key @k@ with the priority @p@. +data Elem a = E + { key :: {-# UNPACK #-} !Key + , prio :: {-# UNPACK #-} !Prio + , value :: a + } deriving (Eq, Show) + +------------------------------------------------------------------------ +-- | A mapping from keys @k@ to priorites @p@. + +data PSQ a = Void + | Winner {-# UNPACK #-} !(Elem a) + !(LTree a) + {-# UNPACK #-} !Key -- max key + deriving (Eq, Show) + +-- | /O(1)/ The number of elements in a queue. +size :: PSQ a -> Int +size Void = 0 +size (Winner _ lt _) = 1 + size' lt + +-- | /O(1)/ True if the queue is empty. +null :: PSQ a -> Bool +null Void = True +null (Winner _ _ _) = False + +-- | /O(log n)/ The priority and value of a given key, or Nothing if +-- the key is not bound. +lookup :: Key -> PSQ a -> Maybe (Prio, a) +lookup k q = case tourView q of + Null -> Nothing + Single (E k' p v) + | k == k' -> Just (p, v) + | otherwise -> Nothing + tl `Play` tr + | k <= maxKey tl -> lookup k tl + | otherwise -> lookup k tr + +------------------------------------------------------------------------ +-- Construction + +empty :: PSQ a +empty = Void + +-- | /O(1)/ Build a queue with one element. +singleton :: Key -> Prio -> a -> PSQ a +singleton k p v = Winner (E k p v) Start k + +------------------------------------------------------------------------ +-- Insertion + +-- | /O(log n)/ Insert a new key, priority and value in the queue. If +-- the key is already present in the queue, the associated priority +-- and value are replaced with the supplied priority and value. +insert :: Key -> Prio -> a -> PSQ a -> PSQ a +insert k p v q = case q of + Void -> singleton k p v + Winner (E k' p' v') Start _ -> case compare k k' of + LT -> singleton k p v `play` singleton k' p' v' + EQ -> singleton k p v + GT -> singleton k' p' v' `play` singleton k p v + Winner e (RLoser _ e' tl m tr) m' + | k <= m -> insert k p v (Winner e tl m) `play` (Winner e' tr m') + | otherwise -> (Winner e tl m) `play` insert k p v (Winner e' tr m') + Winner e (LLoser _ e' tl m tr) m' + | k <= m -> insert k p v (Winner e' tl m) `play` (Winner e tr m') + | otherwise -> (Winner e' tl m) `play` insert k p v (Winner e tr m') + +------------------------------------------------------------------------ +-- Delete/Update + +-- | /O(log n)/ Delete a key and its priority and value from the +-- queue. When the key is not a member of the queue, the original +-- queue is returned. +delete :: Key -> PSQ a -> PSQ a +delete k q = case q of + Void -> empty + Winner (E k' p v) Start _ + | k == k' -> empty + | otherwise -> singleton k' p v + Winner e (RLoser _ e' tl m tr) m' + | k <= m -> delete k (Winner e tl m) `play` (Winner e' tr m') + | otherwise -> (Winner e tl m) `play` delete k (Winner e' tr m') + Winner e (LLoser _ e' tl m tr) m' + | k <= m -> delete k (Winner e' tl m) `play` (Winner e tr m') + | otherwise -> (Winner e' tl m) `play` delete k (Winner e tr m') + +-- | /O(log n)/ Update a priority at a specific key with the result +-- of the provided function. When the key is not a member of the +-- queue, the original queue is returned. +adjust :: (Prio -> Prio) -> Key -> PSQ a -> PSQ a +adjust f k q0 = go q0 + where + go q = case q of + Void -> empty + Winner (E k' p v) Start _ + | k == k' -> singleton k' (f p) v + | otherwise -> singleton k' p v + Winner e (RLoser _ e' tl m tr) m' + | k <= m -> go (Winner e tl m) `unsafePlay` (Winner e' tr m') + | otherwise -> (Winner e tl m) `unsafePlay` go (Winner e' tr m') + Winner e (LLoser _ e' tl m tr) m' + | k <= m -> go (Winner e' tl m) `unsafePlay` (Winner e tr m') + | otherwise -> (Winner e' tl m) `unsafePlay` go (Winner e tr m') +{-# INLINE adjust #-} + +------------------------------------------------------------------------ +-- Conversion + +-- | /O(n*log n)/ Build a queue from a list of key/priority/value +-- tuples. If the list contains more than one priority and value for +-- the same key, the last priority and value for the key is retained. +fromList :: [Elem a] -> PSQ a +fromList = foldr (\(E k p v) q -> insert k p v q) empty + +-- | /O(n)/ Convert to a list of key/priority/value tuples. +toList :: PSQ a -> [Elem a] +toList = toAscList + +-- | /O(n)/ Convert to an ascending list. +toAscList :: PSQ a -> [Elem a] +toAscList q = seqToList (toAscLists q) + +toAscLists :: PSQ a -> Sequ (Elem a) +toAscLists q = case tourView q of + Null -> emptySequ + Single e -> singleSequ e + tl `Play` tr -> toAscLists tl <> toAscLists tr + +-- | /O(n)/ Convert to a descending list. +toDescList :: PSQ a -> [ Elem a ] +toDescList q = seqToList (toDescLists q) + +toDescLists :: PSQ a -> Sequ (Elem a) +toDescLists q = case tourView q of + Null -> emptySequ + Single e -> singleSequ e + tl `Play` tr -> toDescLists tr <> toDescLists tl + +------------------------------------------------------------------------ +-- Min + +-- | /O(1)/ The element with the lowest priority. +findMin :: PSQ a -> Maybe (Elem a) +findMin Void = Nothing +findMin (Winner e _ _) = Just e + +-- | /O(log n)/ Delete the element with the lowest priority. Returns +-- an empty queue if the queue is empty. +deleteMin :: PSQ a -> PSQ a +deleteMin Void = Void +deleteMin (Winner _ t m) = secondBest t m + +-- | /O(log n)/ Retrieve the binding with the least priority, and the +-- rest of the queue stripped of that binding. +minView :: PSQ a -> Maybe (Elem a, PSQ a) +minView Void = Nothing +minView (Winner e t m) = Just (e, secondBest t m) + +secondBest :: LTree a -> Key -> PSQ a +secondBest Start _ = Void +secondBest (LLoser _ e tl m tr) m' = Winner e tl m `play` secondBest tr m' +secondBest (RLoser _ e tl m tr) m' = secondBest tl m `play` Winner e tr m' + +-- | /O(r*(log n - log r))/ Return a list of elements ordered by +-- key whose priorities are at most @pt@. +atMost :: Prio -> PSQ a -> ([Elem a], PSQ a) +atMost pt q = let (sequ, q') = atMosts pt q + in (seqToList sequ, q') + +atMosts :: Prio -> PSQ a -> (Sequ (Elem a), PSQ a) +atMosts !pt q = case q of + (Winner e _ _) + | prio e > pt -> (emptySequ, q) + Void -> (emptySequ, Void) + Winner e Start _ -> (singleSequ e, Void) + Winner e (RLoser _ e' tl m tr) m' -> + let (sequ, q') = atMosts pt (Winner e tl m) + (sequ', q'') = atMosts pt (Winner e' tr m') + in (sequ <> sequ', q' `play` q'') + Winner e (LLoser _ e' tl m tr) m' -> + let (sequ, q') = atMosts pt (Winner e' tl m) + (sequ', q'') = atMosts pt (Winner e tr m') + in (sequ <> sequ', q' `play` q'') + +------------------------------------------------------------------------ +-- Loser tree + +type Size = Int + +data LTree a = Start + | LLoser {-# UNPACK #-} !Size + {-# UNPACK #-} !(Elem a) + !(LTree a) + {-# UNPACK #-} !Key -- split key + !(LTree a) + | RLoser {-# UNPACK #-} !Size + {-# UNPACK #-} !(Elem a) + !(LTree a) + {-# UNPACK #-} !Key -- split key + !(LTree a) + deriving (Eq, Show) + +size' :: LTree a -> Size +size' Start = 0 +size' (LLoser s _ _ _ _) = s +size' (RLoser s _ _ _ _) = s + +left, right :: LTree a -> LTree a + +left Start = moduleError "left" "empty loser tree" +left (LLoser _ _ tl _ _ ) = tl +left (RLoser _ _ tl _ _ ) = tl + +right Start = moduleError "right" "empty loser tree" +right (LLoser _ _ _ _ tr) = tr +right (RLoser _ _ _ _ tr) = tr + +maxKey :: PSQ a -> Key +maxKey Void = moduleError "maxKey" "empty queue" +maxKey (Winner _ _ m) = m + +lloser, rloser :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a +lloser k p v tl m tr = LLoser (1 + size' tl + size' tr) (E k p v) tl m tr +rloser k p v tl m tr = RLoser (1 + size' tl + size' tr) (E k p v) tl m tr + +------------------------------------------------------------------------ +-- Balancing + +-- | Balance factor +omega :: Int +omega = 4 + +lbalance, rbalance :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a + +lbalance k p v l m r + | size' l + size' r < 2 = lloser k p v l m r + | size' r > omega * size' l = lbalanceLeft k p v l m r + | size' l > omega * size' r = lbalanceRight k p v l m r + | otherwise = lloser k p v l m r + +rbalance k p v l m r + | size' l + size' r < 2 = rloser k p v l m r + | size' r > omega * size' l = rbalanceLeft k p v l m r + | size' l > omega * size' r = rbalanceRight k p v l m r + | otherwise = rloser k p v l m r + +lbalanceLeft :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a +lbalanceLeft k p v l m r + | size' (left r) < size' (right r) = lsingleLeft k p v l m r + | otherwise = ldoubleLeft k p v l m r + +lbalanceRight :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a +lbalanceRight k p v l m r + | size' (left l) > size' (right l) = lsingleRight k p v l m r + | otherwise = ldoubleRight k p v l m r + +rbalanceLeft :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a +rbalanceLeft k p v l m r + | size' (left r) < size' (right r) = rsingleLeft k p v l m r + | otherwise = rdoubleLeft k p v l m r + +rbalanceRight :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a +rbalanceRight k p v l m r + | size' (left l) > size' (right l) = rsingleRight k p v l m r + | otherwise = rdoubleRight k p v l m r + +lsingleLeft :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a +lsingleLeft k1 p1 v1 t1 m1 (LLoser _ (E k2 p2 v2) t2 m2 t3) + | p1 <= p2 = lloser k1 p1 v1 (rloser k2 p2 v2 t1 m1 t2) m2 t3 + | otherwise = lloser k2 p2 v2 (lloser k1 p1 v1 t1 m1 t2) m2 t3 +lsingleLeft k1 p1 v1 t1 m1 (RLoser _ (E k2 p2 v2) t2 m2 t3) = + rloser k2 p2 v2 (lloser k1 p1 v1 t1 m1 t2) m2 t3 +lsingleLeft _ _ _ _ _ _ = moduleError "lsingleLeft" "malformed tree" + +rsingleLeft :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a +rsingleLeft k1 p1 v1 t1 m1 (LLoser _ (E k2 p2 v2) t2 m2 t3) = + rloser k1 p1 v1 (rloser k2 p2 v2 t1 m1 t2) m2 t3 +rsingleLeft k1 p1 v1 t1 m1 (RLoser _ (E k2 p2 v2) t2 m2 t3) = + rloser k2 p2 v2 (rloser k1 p1 v1 t1 m1 t2) m2 t3 +rsingleLeft _ _ _ _ _ _ = moduleError "rsingleLeft" "malformed tree" + +lsingleRight :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a +lsingleRight k1 p1 v1 (LLoser _ (E k2 p2 v2) t1 m1 t2) m2 t3 = + lloser k2 p2 v2 t1 m1 (lloser k1 p1 v1 t2 m2 t3) +lsingleRight k1 p1 v1 (RLoser _ (E k2 p2 v2) t1 m1 t2) m2 t3 = + lloser k1 p1 v1 t1 m1 (lloser k2 p2 v2 t2 m2 t3) +lsingleRight _ _ _ _ _ _ = moduleError "lsingleRight" "malformed tree" + +rsingleRight :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a +rsingleRight k1 p1 v1 (LLoser _ (E k2 p2 v2) t1 m1 t2) m2 t3 = + lloser k2 p2 v2 t1 m1 (rloser k1 p1 v1 t2 m2 t3) +rsingleRight k1 p1 v1 (RLoser _ (E k2 p2 v2) t1 m1 t2) m2 t3 + | p1 <= p2 = rloser k1 p1 v1 t1 m1 (lloser k2 p2 v2 t2 m2 t3) + | otherwise = rloser k2 p2 v2 t1 m1 (rloser k1 p1 v1 t2 m2 t3) +rsingleRight _ _ _ _ _ _ = moduleError "rsingleRight" "malformed tree" + +ldoubleLeft :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a +ldoubleLeft k1 p1 v1 t1 m1 (LLoser _ (E k2 p2 v2) t2 m2 t3) = + lsingleLeft k1 p1 v1 t1 m1 (lsingleRight k2 p2 v2 t2 m2 t3) +ldoubleLeft k1 p1 v1 t1 m1 (RLoser _ (E k2 p2 v2) t2 m2 t3) = + lsingleLeft k1 p1 v1 t1 m1 (rsingleRight k2 p2 v2 t2 m2 t3) +ldoubleLeft _ _ _ _ _ _ = moduleError "ldoubleLeft" "malformed tree" + +ldoubleRight :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a +ldoubleRight k1 p1 v1 (LLoser _ (E k2 p2 v2) t1 m1 t2) m2 t3 = + lsingleRight k1 p1 v1 (lsingleLeft k2 p2 v2 t1 m1 t2) m2 t3 +ldoubleRight k1 p1 v1 (RLoser _ (E k2 p2 v2) t1 m1 t2) m2 t3 = + lsingleRight k1 p1 v1 (rsingleLeft k2 p2 v2 t1 m1 t2) m2 t3 +ldoubleRight _ _ _ _ _ _ = moduleError "ldoubleRight" "malformed tree" + +rdoubleLeft :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a +rdoubleLeft k1 p1 v1 t1 m1 (LLoser _ (E k2 p2 v2) t2 m2 t3) = + rsingleLeft k1 p1 v1 t1 m1 (lsingleRight k2 p2 v2 t2 m2 t3) +rdoubleLeft k1 p1 v1 t1 m1 (RLoser _ (E k2 p2 v2) t2 m2 t3) = + rsingleLeft k1 p1 v1 t1 m1 (rsingleRight k2 p2 v2 t2 m2 t3) +rdoubleLeft _ _ _ _ _ _ = moduleError "rdoubleLeft" "malformed tree" + +rdoubleRight :: Key -> Prio -> a -> LTree a -> Key -> LTree a -> LTree a +rdoubleRight k1 p1 v1 (LLoser _ (E k2 p2 v2) t1 m1 t2) m2 t3 = + rsingleRight k1 p1 v1 (lsingleLeft k2 p2 v2 t1 m1 t2) m2 t3 +rdoubleRight k1 p1 v1 (RLoser _ (E k2 p2 v2) t1 m1 t2) m2 t3 = + rsingleRight k1 p1 v1 (rsingleLeft k2 p2 v2 t1 m1 t2) m2 t3 +rdoubleRight _ _ _ _ _ _ = moduleError "rdoubleRight" "malformed tree" + +-- | Take two pennants and returns a new pennant that is the union of +-- the two with the precondition that the keys in the first tree are +-- strictly smaller than the keys in the second tree. +play :: PSQ a -> PSQ a -> PSQ a +Void `play` t' = t' +t `play` Void = t +Winner e@(E k p v) t m `play` Winner e'@(E k' p' v') t' m' + | p <= p' = Winner e (rbalance k' p' v' t m t') m' + | otherwise = Winner e' (lbalance k p v t m t') m' +{-# INLINE play #-} + +-- | A version of 'play' that can be used if the shape of the tree has +-- not changed or if the tree is known to be balanced. +unsafePlay :: PSQ a -> PSQ a -> PSQ a +Void `unsafePlay` t' = t' +t `unsafePlay` Void = t +Winner e@(E k p v) t m `unsafePlay` Winner e'@(E k' p' v') t' m' + | p <= p' = Winner e (rloser k' p' v' t m t') m' + | otherwise = Winner e' (lloser k p v t m t') m' +{-# INLINE unsafePlay #-} + +data TourView a = Null + | Single {-# UNPACK #-} !(Elem a) + | (PSQ a) `Play` (PSQ a) + +tourView :: PSQ a -> TourView a +tourView Void = Null +tourView (Winner e Start _) = Single e +tourView (Winner e (RLoser _ e' tl m tr) m') = + Winner e tl m `Play` Winner e' tr m' +tourView (Winner e (LLoser _ e' tl m tr) m') = + Winner e' tl m `Play` Winner e tr m' + +------------------------------------------------------------------------ +-- Utility functions + +moduleError :: String -> String -> a +moduleError fun msg = error ("GHC.Event.PSQ." ++ fun ++ ':' : ' ' : msg) +{-# NOINLINE moduleError #-} + +------------------------------------------------------------------------ +-- Hughes's efficient sequence type + +newtype Sequ a = Sequ ([a] -> [a]) + +emptySequ :: Sequ a +emptySequ = Sequ (\as -> as) + +singleSequ :: a -> Sequ a +singleSequ a = Sequ (\as -> a : as) + +(<>) :: Sequ a -> Sequ a -> Sequ a +Sequ x1 <> Sequ x2 = Sequ (\as -> x1 (x2 as)) +infixr 5 <> + +seqToList :: Sequ a -> [a] +seqToList (Sequ x) = x [] + +instance Show a => Show (Sequ a) where + showsPrec d a = showsPrec d (seqToList a) + diff --git a/Network/DNS/Cache/Sync.hs b/Network/DNS/Cache/Sync.hs new file mode 100644 --- /dev/null +++ b/Network/DNS/Cache/Sync.hs @@ -0,0 +1,78 @@ +{-# LANGUAGE BangPatterns #-} + +module Network.DNS.Cache.Sync ( + ConcVar + , newConcVar + , wait + , waitIncrease + , decrease + , ActiveVar + , newActiveVar + , tell + , listen + , ActiveRef + , newActiveRef + , lookupActiveRef + , insertActiveRef + , deleteActiveRef + ) where + +import Control.Applicative ((<$>)) +import Control.Concurrent.STM +import Data.Map (Map) +import qualified Data.Map as Map +import Network.DNS.Cache.Types +import Data.IORef (newIORef, readIORef, atomicModifyIORef', IORef) + +---------------------------------------------------------------- + +newtype ConcVar = ConcVar (TVar Int) + +newConcVar :: IO ConcVar +newConcVar = ConcVar <$> newTVarIO 0 + +wait :: ConcVar -> (Int -> Bool) -> IO () +wait (ConcVar var) cond = atomically $ do + x <- readTVar var + check (cond x) + +waitIncrease :: ConcVar -> Int -> IO () +waitIncrease (ConcVar var) lim = atomically $ do + x <- readTVar var + check (x < lim) + let !x' = x + 1 + writeTVar var x' + +decrease :: ConcVar -> IO () +decrease (ConcVar var) = atomically $ modifyTVar' var (subtract 1) + +---------------------------------------------------------------- + +newtype ActiveVar = ActiveVar (TMVar (Either DNSError Result)) + +newActiveVar :: IO ActiveVar +newActiveVar = ActiveVar <$> newEmptyTMVarIO + +tell :: ActiveVar -> Either DNSError Result -> IO () +tell (ActiveVar var) r = atomically $ putTMVar var r + +listen :: ActiveVar -> IO (Either DNSError Result) +listen (ActiveVar var) = atomically $ readTMVar var + +---------------------------------------------------------------- + +newtype ActiveRef = ActiveRef (IORef (Map Key ActiveVar)) + +newActiveRef :: IO ActiveRef +newActiveRef = ActiveRef <$> newIORef Map.empty + +lookupActiveRef :: Key -> ActiveRef -> IO (Maybe ActiveVar) +lookupActiveRef key (ActiveRef ref) = Map.lookup key <$> readIORef ref + +insertActiveRef :: Key -> ActiveVar -> ActiveRef -> IO () +insertActiveRef key avar (ActiveRef ref) = + atomicModifyIORef' ref $ \mp -> (Map.insert key avar mp, ()) + +deleteActiveRef :: Key -> ActiveRef -> IO () +deleteActiveRef key (ActiveRef ref) = + atomicModifyIORef' ref $ \mp -> (Map.delete key mp, ()) diff --git a/Network/DNS/Cache/Types.hs b/Network/DNS/Cache/Types.hs new file mode 100644 --- /dev/null +++ b/Network/DNS/Cache/Types.hs @@ -0,0 +1,52 @@ +module Network.DNS.Cache.Types ( + Hash + , Key(..) + , Prio + , Value(..) + , Entry + , Result(..) + , TTL + , HostAddress + , Domain + , DNSError(..) + ) where + +import Data.Array.Unboxed (UArray) +import Data.ByteString.Short (ShortByteString) +import Data.IORef (IORef) +import Data.Time (UTCTime) +import Network.DNS (Domain, DNSError(..)) +import Network.Socket (HostAddress) + +type Hash = Int + +data Key = Key !Hash -- making lookup faster + !ShortByteString -- avoiding memory fragmentation + deriving (Ord,Show) + +instance Eq Key where + -- Just ensuring the order of evaluation. + Key h1 k1 == Key h2 k2 = h1 == h2 && k1 == k2 + +type Prio = UTCTime + +-- fixme: if UArray causes memory fragments, +-- we should use real-time queue instaed. +data Value = Value (UArray Int HostAddress) (IORef Int) + +instance Show Value where + show (Value a _) = show a + +type TTL = Int + +-- | Information of positive result. +data Result = + -- | An address obtained from the cache. + Hit HostAddress + -- | An address resolved from cache DNS servers. + | Resolved HostAddress + -- | Specified domain is IP address. So, it is converted into a numeric address. + | Numeric HostAddress + deriving (Eq,Show) + +type Entry = Either DNSError Value diff --git a/Network/DNS/Cache/Utils.hs b/Network/DNS/Cache/Utils.hs new file mode 100644 --- /dev/null +++ b/Network/DNS/Cache/Utils.hs @@ -0,0 +1,22 @@ +module Network.DNS.Cache.Utils where + +import qualified Data.ByteString.Char8 as BS +import Data.Char (isDigit) +import Network.DNS.Cache.Types + +isIPAddr :: Domain -> Bool +isIPAddr hn = length groups == 4 && all ip groups + where + groups = BS.split '.' hn + ip x = BS.length x <= 3 + && BS.all isDigit x + && read (BS.unpack x) <= (255 :: Int) + +fromResult :: Result -> HostAddress +fromResult (Hit addr) = addr +fromResult (Resolved addr) = addr +fromResult (Numeric addr) = addr + +fromEither :: Either DNSError Result -> Maybe HostAddress +fromEither (Right res) = Just (fromResult res) +fromEither (Left _) = Nothing diff --git a/Network/DNS/Cache/Value.hs b/Network/DNS/Cache/Value.hs new file mode 100644 --- /dev/null +++ b/Network/DNS/Cache/Value.hs @@ -0,0 +1,28 @@ +{-# LANGUAGE BangPatterns #-} + +module Network.DNS.Cache.Value where + +import Data.Array.Unboxed +import Data.IORef (newIORef, atomicModifyIORef') +import Network.DNS.Cache.Types + +positiveEntry :: [HostAddress] -> IO Entry +positiveEntry addrs = do + ref <- newIORef next + let !val = Value arr ref + return $! Right val + where + !siz = length addrs + !next = adjust 0 siz + !arr = listArray (0,siz-1) addrs + +rotate :: Value -> IO HostAddress +rotate (Value a ref) = do + let (_, siz) = bounds a + j <- atomicModifyIORef' ref $ \i -> (adjust i siz, i) + let !addr = a ! j + return addr + +adjust :: Int -> Int -> Int +adjust i 0 = i +adjust i n = let !x = (i + 1) `mod` n in x 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/concurrent-dns-cache.cabal b/concurrent-dns-cache.cabal new file mode 100644 --- /dev/null +++ b/concurrent-dns-cache.cabal @@ -0,0 +1,60 @@ +Name: concurrent-dns-cache +Version: 0.0.0 +Author: Kazu Yamamoto +Maintainer: Kazu Yamamoto +License: BSD3 +License-File: LICENSE +Synopsis: Concurrent DNS cache +Description: Concurrent DNS cache +Category: Network +Cabal-Version: >= 1.10 +Build-Type: Simple + +Library + Default-Language: Haskell2010 + GHC-Options: -Wall + Exposed-Modules: Network.DNS.Cache + Other-Modules: Network.DNS.Cache.Cache + Network.DNS.Cache.PSQ + Network.DNS.Cache.Sync + Network.DNS.Cache.Types + Network.DNS.Cache.Utils + Network.DNS.Cache.Value + Build-Depends: base >= 4 && < 5 + , array + , async + , bytestring >= 0.10.4.0 + , containers + , dns + , hashable + , iproute + , network + , stm + , time + +Executable main + Main-Is: main.hs + Default-Language: Haskell2010 + GHC-Options: -Wall -threaded + Build-Depends: base >= 4 && < 5 + , array + , async + , bytestring >= 0.10.4.0 + , containers + , dns + , hashable + , iproute + , network + , stm + , time + +Test-Suite spec + Main-Is: Spec.hs + Default-Language: Haskell2010 + Other-modules: CacheSpec + Hs-Source-Dirs: test + Type: exitcode-stdio-1.0 + Build-Depends: base >= 4 && < 5 + , dns + , concurrent-dns-cache + , hspec diff --git a/main.hs b/main.hs new file mode 100644 --- /dev/null +++ b/main.hs @@ -0,0 +1,53 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Main where + +import Control.Concurrent (forkIO) +import Control.Exception (try, SomeException(..)) +import Control.Monad (void, when) +import qualified Data.ByteString.Char8 as BS +import Network.DNS +import Network.DNS.Cache as DNSC +import Data.Time + +confs :: [ResolvConf] +confs = [ + defaultResolvConf { resolvInfo = RCHostName "8.8.8.8" } + , defaultResolvConf { resolvInfo = RCHostName "8.8.4.4" } + ] + +maxCon :: Int +maxCon = 50 + +cacheConf :: DNSCacheConf +cacheConf = DNSCacheConf { + resolvConfs = confs + , maxConcurrency = maxCon + , minTTL = 60 + , maxTTL = 300 + , negativeTTL = 300 + } + +main :: IO () +main = do + beg <- getCurrentTime + withDNSCache cacheConf (loop 1 beg) + where + loop :: Int -> UTCTime -> DNSCache -> IO () + loop n beg cache = do + when (n `mod` 1000 == 0) $ do + cur <- getCurrentTime + putStrLn $ show n ++ ": " ++ show (cur `diffUTCTime` beg) + edom <- try BS.getLine + case edom of + Left (SomeException _) -> do + wait cache (== 0) + putStrLn "Done." + Right dom -> do + wait cache (< maxCon) + void $ forkIO (DNSC.resolve cache dom >>= p dom) + loop (n+1) beg cache + p _ (Right _) = return () + p dom (Left e) = do + putStr $ show e ++ " " + BS.putStrLn dom diff --git a/test/CacheSpec.hs b/test/CacheSpec.hs new file mode 100644 --- /dev/null +++ b/test/CacheSpec.hs @@ -0,0 +1,44 @@ +{-# LANGUAGE OverloadedStrings #-} + +module CacheSpec where + +import Test.Hspec +import Network.DNS hiding (lookup) +import Network.DNS.Cache +import Prelude hiding (lookup) + +cacheConf :: DNSCacheConf +cacheConf = DNSCacheConf { + resolvConfs = [ + defaultResolvConf { resolvInfo = RCHostName "8.8.8.8" } + , defaultResolvConf { resolvInfo = RCHostName "8.8.4.4" } + ] + , maxConcurrency = 10 + , minTTL = 60 + , maxTTL = 300 + , negativeTTL = 300 + } + +spec :: Spec +spec = describe "withDnsCache" $ do + it "resolves domains and caches addresses" $ withDNSCache cacheConf $ \cache -> do + let dom = "www.example.com" + addr = 2010691677 + resolve cache dom `shouldReturn` Right (Resolved addr) + resolveCache cache dom `shouldReturn` Just (Right (Hit addr)) + lookupCache cache dom `shouldReturn` Just addr + lookup cache dom `shouldReturn` Just addr + it "resolves domains and caches nagative" $ withDNSCache cacheConf $ \cache -> do + let dom = "not-exist.com" + err = UnexpectedRDATA + resolve cache dom `shouldReturn` Left err + resolveCache cache dom `shouldReturn` Just (Left err) + lookupCache cache dom `shouldReturn` Nothing + lookup cache dom `shouldReturn` Nothing + it "resolves domains and caches nagative" $ withDNSCache cacheConf $ \cache -> do + let dom = "non-exist.org" + err = NameError + resolve cache dom `shouldReturn` Left err + resolveCache cache dom `shouldReturn` Just (Left err) + lookupCache cache dom `shouldReturn` Nothing + lookup cache dom `shouldReturn` Nothing diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,1 @@ +{-# OPTIONS_GHC -F -pgmF hspec-discover #-}