Skip to content

HTTPS clone URL

Subversion checkout URL

You can clone with HTTPS or Subversion.

Download ZIP
branch: strict
Fetching contributors…

Cannot retrieve contributors at this time

516 lines (445 sloc) 19.092 kb
{-# LANGUAGE OverloadedStrings, GADTs, KindSignatures, TypeFamilies, BangPatterns #-}
import Blaze.ByteString.Builder
import Data.Word
import Data.Int
import Data.Bits
import qualified Data.ByteString.Char8 as B8
import qualified Data.ByteString.Lazy.Char8 as BL8
import Data.Monoid
import Data.Serialize
import Data.Serialize.Get
import Control.Monad
import Control.Monad.IO.Class
import Control.Applicative
import Debug.Trace
import qualified System.IO.Unsafe as IOU
import Data.Conduit hiding (Done,sequence)
import qualified Data.Conduit.Binary as C
import qualified Database.Redis as R
redisHeader = BL8.pack "REDIS0003"
-- Redis types
typeString = 0x00
typeList = 0x01
typeSet = 0x02
typeZset = 0x03
typeHash = 0x04
-- Encoded types
typeHashZipmap = 0x09
typeListZiplist = 0x0a
typeSetIntset = 0x0b
typeZsetZiplist = 0x0c
-- Length encodings for types
encInt8 = 0x00
encInt16 = 0x01
encInt32 = 0x02
encLzf = 0x03
-- Redis Opcodes
opcodeEof = 0xff
opcodeSelectdb = 0xfe
opcodeExpiretime = 0xfd
opcodeExpiretimems = 0xfc
-- Redis Length Codes
len6bit = 0x00
len14bit = 0x01
len32bit = 0x02
len_enc = 0x04
data RDBObj = RDBString B8.ByteString |
RDBList [B8.ByteString] |
RDBSet [B8.ByteString] |
RDBZSet [(B8.ByteString,Double)] |
RDBHash [(B8.ByteString,B8.ByteString)] |
RDBPair (Maybe Integer,B8.ByteString,RDBObj) |
RDBDatabase Integer [RDBObj] |
RDBSelect Integer |
RDBNull |
RDB [RDBObj] deriving (Show,Eq)
toZsetPairs :: [B8.ByteString] -> [(B8.ByteString,Double)]
toZsetPairs [] = []
toZsetPairs l = (x, read $ B8.unpack y) : toZsetPairs xs where
(x:(y:[]),xs) = splitAt 2 l
getEncoding :: Word8 -> Word8
getEncoding = flip shift (-6) . (.&.) 0xC0
getSecondEncoding :: Word8 -> Word8
getSecondEncoding = flip shift (-4) . (.&.) 0x30
get6bitLen :: Word8 -> Integer
get6bitLen = fromIntegral . (.&.) 0x3f
get14bitLen :: Word16 -> Word16 -> Integer
get14bitLen f s = fromIntegral (shift (f .&. 0x3f) 8 .|. s)
combineDistances :: Word16 -> Word16 -> Word16
combineDistances h l = shift h 8 .|. l
loadTimeMs :: Get Integer
loadTimeMs = do
time <- getWord64host
return $ fromIntegral (fromIntegral time :: Int64)
loadTime :: Get Integer
loadTime = do
time <- getWord32host
return $ fromIntegral (fromIntegral time :: Int64)
loadLen :: Get (Bool, Integer)
loadLen = do
first <- getWord8
case getEncoding first of
0x00 -> return (False, get6bitLen first)
0x01 -> do
second <- getWord8
return (False, get14bitLen (fromIntegral first) (fromIntegral second) )
0x02 -> do
len <- getWord32be
return (False, fromIntegral len)
0x03 -> return (True, get6bitLen first)
loadIntegerObj :: Integer -> Bool -> Get B8.ByteString
loadIntegerObj len enc = case len of
0x00 -> do
str <- getWord8
return $ B8.pack $ show (fromIntegral str :: Int8)
0x01 -> do
str <- getWord16le
return $ B8.pack $ show (fromIntegral str :: Int16)
0x02 -> do
str <- getWord32le
return $ B8.pack $ show (fromIntegral str :: Int32)
loadDoubleValue :: Get Double
loadDoubleValue = do
len <- getWord8
case len of
0xff -> return $ read "-Infinity"
0xfe -> return $ read "Infinity"
0xfd -> return $ read "NaN"
l -> do
val <- getByteString (fromIntegral l)
return $ read $ B8.unpack val
loadLzfStr :: Get B8.ByteString
loadLzfStr = do
(clenEnc, !clen) <- loadLen
(lenEnc, !len) <- loadLen
!str <- getLazyByteString (fromIntegral clen)
return $ decompressLzfStr str
decompressLzfStr :: BL8.ByteString -> B8.ByteString
decompressLzfStr s = B8.concat $ BL8.toChunks str where
(Right !str) = runGetLazy (parseLzf BL8.empty) s
parseLzf :: BL8.ByteString -> Get BL8.ByteString
parseLzf decodedString = do
empty <- isEmpty
if empty then return BL8.empty else
do
h <- getWord8
let len = shift (h .&. 0xe0) (-5)
let high_d = h .&. 0x1f
obj <- case len of
0 -> getLazyByteString ((fromIntegral high_d :: Int64) + 1)
7 -> do
new_len <- getWord8
let full_len = (fromIntegral new_len :: Word16) + 7 + 2
low_d <- getWord8
let distance = combineDistances (fromIntegral high_d) (fromIntegral low_d)
return $ repCopy decodedString distance full_len
l -> do
low_d <- getWord8
let distance = combineDistances (fromIntegral high_d) (fromIntegral low_d)
return $ repCopy decodedString distance ((fromIntegral l :: Word16) + 2)
rest <- parseLzf (BL8.append decodedString obj)
return (BL8.append obj rest)
repCopy :: (Integral a) => BL8.ByteString -> a -> a -> BL8.ByteString
repCopy str dist len = BL8.take (fromIntegral len) (BL8.cycle window) where
window = BL8.drop (olen - fromIntegral dist - 1) str
olen = BL8.length str
loadStringObj :: Bool -> Get B8.ByteString
loadStringObj enc = do
(isEncType,len) <- loadLen
if isEncType
then case len of
0x00 -> loadIntegerObj len enc
0x01 -> loadIntegerObj len enc
0x02 -> loadIntegerObj len enc
0x03 -> loadLzfStr
else getByteString (fromIntegral len)
loadListObj :: Get [B8.ByteString]
loadListObj = do
(isEncType,len) <- loadLen
replicateM (fromIntegral len) (loadStringObj True)
loadZSetPair :: Get (B8.ByteString, Double)
loadZSetPair = do
value <- loadStringObj True
weight <- loadDoubleValue
return (value,weight)
loadZSetObj :: Get [(B8.ByteString,Double)]
loadZSetObj = do
(isEncType,len) <- loadLen
replicateM (fromIntegral len) loadZSetPair
loadZipListObj :: Get [B8.ByteString]
loadZipListObj = do
ziplen <- getWord32le
offset <- getWord32le
num_entries <- getWord16le
obj <- loadZipListMembers
eoz <- getWord8
return obj
loadZipListMembers :: Get [B8.ByteString]
loadZipListMembers = do
opc <- lookAhead getWord8
if opc == 0xff
then return []
else do
obj <- getZipListMember
rest <- loadZipListMembers
return (obj:rest)
getZipListMember :: Get B8.ByteString
getZipListMember = do
prevLen <- getZipLen
header <- getWord8
case (getEncoding header, getSecondEncoding header) of
(0x00,_) -> do
let len = get6bitLen header
getByteString (fromIntegral len)
(0x01,_) -> do
second_part <- getWord8
let len = get14bitLen (fromIntegral header) (fromIntegral second_part)
getByteString (fromIntegral len)
(0x02,_) -> do
len <- getWord32be
getByteString (fromIntegral len)
(0x03,0x00) -> do
obj <- getWord16le
return $ B8.pack $ show (fromIntegral obj :: Int16)
(0x03,0x01) -> do
obj <- getWord32le
return $ B8.pack $ show (fromIntegral obj :: Int32)
(0x03,0x02) -> do
obj <- getWord64le
return $ B8.pack $ show (fromIntegral obj :: Int64)
getZipLen :: Get Integer
getZipLen = do
prevLen <- getWord8
case prevLen of
0xfe -> do
len <- getWord32le
return (fromIntegral len)
_ -> return (fromIntegral prevLen)
loadIntSetObj :: Get [B8.ByteString]
loadIntSetObj = do
enc <- getWord32le
setlen <- getWord32le
replicateM (fromIntegral setlen) (loadIntSetMember enc)
loadIntSetMember :: Word32 -> Get B8.ByteString
loadIntSetMember enc = case enc of
0x02 -> do
obj <- getWord16le
return $ B8.pack $ show (fromIntegral obj :: Int16)
0x04 -> do
obj <- getWord32le
return $ B8.pack $ show (fromIntegral obj :: Int32)
0x08 -> do
obj <- getWord64le
return $ B8.pack $ show (fromIntegral obj :: Int64)
l -> do
obj <- trace (show enc) $ getWord8
return "1"
loadHashObj :: Get [(B8.ByteString,B8.ByteString)]
loadHashObj = do
(isEncType,len) <- loadLen
replicateM (fromIntegral len) loadHashMember
loadHashMember :: Get (B8.ByteString,B8.ByteString)
loadHashMember = do
key <- loadStringObj True
value <- loadStringObj True
return (key,value)
loadZipMapObj :: Get [(B8.ByteString,B8.ByteString)]
loadZipMapObj = do
zmlen <- getWord8
loadZipMapMembers
loadZipMapMembers :: Get [(B8.ByteString,B8.ByteString)]
loadZipMapMembers = do
is_eoz <- lookAhead getWord8
case is_eoz of
0xff -> do
skip 1
return []
_ -> do
obj <- loadZipMapMember
rest <- loadZipMapMembers
return (obj:rest)
loadZipMapMember :: Get (B8.ByteString,B8.ByteString)
loadZipMapMember = do
first_len <- getWord8
key_len <- getZipMapMemberLen first_len
key <- getByteString (fromIntegral key_len)
first_len_v <- getWord8
val_len <- getZipMapMemberLen first_len_v
free <- getWord8
val <- getByteString val_len
extra <- getByteString (fromIntegral free)
return (key,val)
getZipMapMemberLen :: Word8 -> Get Int
getZipMapMemberLen first_len = case first_len of
0xfd -> do
l <- getWord32host
return (fromIntegral l)
0xfe -> return 0xfe
l -> return (fromIntegral l)
loadObj :: Word8 -> Get RDBObj
loadObj t = case t of
-- ^ Load a string value
0x00 -> do
obj <- loadStringObj True
return (RDBString obj)
-- ^ Load a list value
0x01 -> do
obj <- loadListObj
return (RDBList obj)
-- ^ Load a set value
0x02 -> do
obj <- loadListObj
return (RDBSet obj)
-- ^ Load a sorted set value
0x03 -> do
obj <- loadZSetObj
return (RDBZSet obj)
-- ^ Load a hash value
0x04 -> do
obj <- loadHashObj
return (RDBHash obj)
-- ^ Load a zipmap encoded hash
0x09 -> do
binstr <- loadStringObj True
let (Right obj) = runGet loadZipMapObj binstr
return (RDBHash obj)
-- ^ Load a ziplist encoded list
0x0a -> do
binstr <- loadStringObj True
let (Right obj) = runGet loadZipListObj binstr
return (RDBList obj)
-- ^ Load an intset encoded set
0x0b -> do
binstr <- loadStringObj True
let (Right obj) = runGet loadIntSetObj binstr
return (RDBSet obj)
-- ^ Load a ziplist encoded zset
0x0c -> do
binstr <- loadStringObj True
let (Right obj) = runGet loadZipListObj binstr
return (RDBZSet $ toZsetPairs obj)
loadObjs :: Get [RDBObj]
loadObjs = do
code <- lookAhead getWord8
case code of
0xfd -> do
skip 1
expire <- loadTime
getPairs (Just expire)
0xfc -> do
skip 1
expire <- loadTimeMs
getPairs (Just expire)
0xfe -> return []
0xff -> return []
_ -> getPairs Nothing
getPairs :: Maybe Integer -> Get [RDBObj]
getPairs ex = do
t <- getWord8
key <- loadStringObj False
obj <- loadObj t
rest <- loadObjs
return (RDBPair (ex,key,obj):rest)
getPair :: Maybe Integer -> Get RDBObj
getPair ex = do
t <- getWord8
key <- loadStringObj False
obj <- loadObj t
return $ RDBPair (ex,key,obj)
getDBs :: Get [RDBObj]
getDBs = do
opc <- lookAhead getWord8
if opc == opcodeSelectdb
then do
skip 1
(isEncType,dbnum) <- loadLen
objs <- loadObjs
rest <- getDBs
return (objs ++ rest)
else return []
getObjInc :: Get RDBObj
getObjInc = do
opc <- lookAhead getWord8
if opc == opcodeSelectdb
then do
skip 1
(isEncType,dbnum) <- loadLen
return $ RDBSelect (fromIntegral dbnum)
else do
code <- lookAhead getWord8
case code of
0xfd -> do
skip 1
expire <- loadTime
getPair (Just expire)
0xfc -> do
skip 1
expire <- loadTimeMs
getPair (Just expire)
0xfe -> do
return RDBNull
0xff -> do
skip 1
return RDBNull
_ -> getPair Nothing
{-main = do-}
{-testf <- BL8.readFile "./dump.rdb"-}
{-print $ show (decode testf :: RDBObj)-}
{-main = do-}
{-testf <- BL8.readFile "./dump.rdb"-}
{-runGet (processRDB_ printRDBObj) testf-}
repParse !input (!st,Nothing) = case result of
(Partial parser) -> (st,Just parser)
(Done !res "") -> (res:st,Just (\input -> Done RDBNull input))
(Done !res !leftover) -> repParse leftover (res:st,Nothing)
where
!result = runGetPartial getObjInc input
repParse !input (!st,Just parser) = case result of
(Partial parser) -> (st,Just parser)
(Done !res "") -> (res:st,Just (\input -> Done RDBNull input))
(Done !res !leftover) -> repParse leftover (res:st,Nothing)
where
!result = parser input
printRDB :: ResourceIO m => Sink B8.ByteString m ()
printRDB =
sinkState Nothing
pushRDB
(\state -> return ())
pushRDB Nothing !input = do liftIO $ mapM_ (\x -> if x == RDBNull then return () else print x) st
return $ StateProcessing p
where
(Done r l) = runGetPartial (getBytes 9) input
(!st,!p) = repParse l ([],Nothing)
pushRDB (Just parser) !input = do liftIO $ mapM_ (\x -> if x == RDBNull then return () else print x) st
return $ StateProcessing p
where
(!st,!p) = repParse input ([],Just parser)
loadRDB :: ResourceIO m => R.Connection -> Sink B8.ByteString m (R.Redis (Either R.Reply R.Status))
loadRDB c =
sinkState Nothing
(pushLoad c)
(\state -> return (R.ping))
saveObj :: RDBObj -> R.Redis ()
saveObj RDBNull = do R.ping >> return ()
saveObj (RDBSelect i) = R.select i >> return ()
saveObj (RDBPair (exp,key,RDBString val)) = R.set key val >> setExpire exp key
saveObj (RDBPair (exp,key,RDBList vals)) = R.rpush key vals >> setExpire exp key
saveObj (RDBPair (exp,key,RDBSet vals)) = R.sadd key vals >> setExpire exp key
saveObj (RDBPair (exp,key,RDBZSet vals)) = (R.zadd key $ map (\(x,y) -> (y,x)) vals) >> setExpire exp key
saveObj (RDBPair (exp,key,RDBHash vals)) = R.hmset key vals >> setExpire exp key
setExpire Nothing k = return ()
setExpire (Just exp) k = R.expireat k exp >> return ()
pushLoad c Nothing !input = do liftIO $ R.runRedis c $ sequence_ $ reverse $ map saveObj st
return $ StateProcessing p
where
(Done r l) = runGetPartial (getBytes 9) input
(!st,!p) = repParse l ([],Nothing)
pushLoad c (Just parser) !input = do liftIO $ R.runRedis c $ sequence $ reverse $ map saveObj st
return $ StateProcessing p
where
(!st,!p) = repParse input ([],Just parser)
{-main = do-}
{-runResourceT $ C.sourceFile "./dump.rdb" $$ printRDB-}
main = do
conn <- R.connect R.defaultConnectInfo
runResourceT $ C.sourceFile "./source.rdb" $$ (loadRDB conn)
Jump to Line
Something went wrong with that request. Please try again.