diff --git a/BENCHMARKS.md b/BENCHMARKS.md index 321c149..40c3616 100644 --- a/BENCHMARKS.md +++ b/BENCHMARKS.md @@ -28,31 +28,36 @@ The second column is wall clock time in seconds, the third is peak heap memory a ### Materializing 100_000 rows with 13 columns each into a List of Records This runs with 2 concurrent queries, 10 times over: + +This benchmark is unfair towards both hpgsql and postgresql-simple because the row decoder is Generically derived for them while it is hand-written for hasql. ```csv -postgresql-simple Record List (100000 rows),11.90,142.65M,101.6 -hasql Record List (100000 rows),6.279,142.48M,78.0 -hpgsql Record List (100000 rows),3.886,72.07M,98.2 +postgresql-simple Record List (100000 rows),12.10,142.65M,120.5 +hasql Record List (100000 rows),6.314,142.48M,77.8 +hpgsql Record List (100000 rows),4.092,72.07M,119.7 ``` ### Materializing 100_000 rows with 13 columns each into a List of Tuples This runs with 2 concurrent queries, 10 times over: ```csv -postgresql-simple Tuple List (100000 rows),15.03,142.78M,149.8 -hasql Tuple List (100000 rows),8.376,142.48M,195.4 -hpgsql Tuple List (100000 rows),4.689,72.07M,149.4 +postgresql-simple Tuple List (100000 rows),14.74,142.52M,144.9 +hasql Tuple List (100000 rows),8.263,142.48M,202.2 +hpgsql Tuple List (100000 rows),4.648,72.07M,150.4 ``` ### Streaming 100_000 rows with 13 columns as Records This runs with 2 concurrent queries, 10 times over. -Hpgsql's implementation streams directly from the socket while the others use cursors, so + +This benchmark is unfair towards both hpgsql and postgresql-simple because the row decoder is Generically derived for them while it is hand-written for hasql. + +However, Hpgsql's implementation streams directly from the socket while the others use cursors, so it might not be a fair comparison in terms of implementation (e.g. you can advance multiple cursors simultaneously, but not hpgsql's Streamed-from-socket streams). ```csv -streaming-postgresql-simple Record Stream (100000 rows),16.69,73.24M,0.1 -postgresql-simple Record fold (100000 rows),13.39,78.29M,0.2 -hpgsql Record Stream (100000 rows),1.457,72.07M,0.2 +streaming-postgresql-simple Record Stream (100000 rows),13.32,73.41M,0.0 +postgresql-simple Record fold (100000 rows),13.47,77.81M,0.0 +hpgsql Record Stream (100000 rows),1.421,72.07M,0.0 ``` ### Streaming 100_000 rows with 13 columns as Tuples @@ -62,9 +67,9 @@ Hpgsql's implementation streams directly from the socket while the others use cu it might not be a fair comparison in terms of implementation (e.g. you can advance multiple cursors simultaneously, but not hpgsql's Streamed-from-socket streams). ```csv -streaming-postgresql-simple Tuple Stream (100000 rows),13.89,73.28M,0.1 -postgresql-simple Tuple fold (100000 rows),13.69,81.99M,0.2 -hpgsql Tuple Stream (100000 rows),1.076,72.07M,0.2 +streaming-postgresql-simple Tuple Stream (100000 rows),14.10,73.26M,0.0 +postgresql-simple Tuple fold (100000 rows),13.45,84.42M,0.0 +hpgsql Tuple Stream (100000 rows),1.025,72.07M,0.0 ``` ### COPY FROM STDIN @@ -72,6 +77,6 @@ hpgsql Tuple Stream (100000 rows),1.076,72.07M,0.2 This compares hpgsql's binary copy to a `forM` loop writing text rows. ```csv -postgresql-simple text COPY (100000 rows),1.353,72.10M,3.9 -hpgsql copyFromS binary COPY (100000 rows),1.239,72.07M,11.0 +postgresql-simple text COPY (100000 rows),1.348,72.10M,3.8 +hpgsql copyFromS binary COPY (100000 rows),672.1,72.07M,10.8 ``` diff --git a/README.md b/README.md index 1fde719..03ecaa5 100644 --- a/README.md +++ b/README.md @@ -51,11 +51,11 @@ You should start by swapping all of "postgresql-simple", "postgresql-libpq", and ## Performance -Some benchmarks show materializing large query results with hpgsql takes 31-33% the time postgresql-simple takes, and 56-62% the time hasql takes (on my computer, Linux x64, GHC 9.10.3, compiled with -O1). +Some benchmarks show materializing large query results with hpgsql takes 31-34% the time postgresql-simple takes, and 55-65% the time hasql takes (on my computer, Linux x64, GHC 9.10.3, compiled with -O1). When comparing hpgsql's Stream querying, hpgsql takes 9-11% the time of both [streaming-postgresql-simple](https://hackage.haskell.org/package/streaming-postgresql-simple) and postgresql-simple's cursor folding functions, although this might not be a fair comparison for some use cases. -hpgsql's binary COPY runs in about 92% the time of postgresql-simple's textual COPY. +hpgsql's binary COPY runs in ~50% the time of postgresql-simple's textual COPY. Peak allocated memory is harder to analyze. diff --git a/hpgsql-benchmarks/src/Main.hs b/hpgsql-benchmarks/src/Main.hs index 0bf4500..081cc9f 100644 --- a/hpgsql-benchmarks/src/Main.hs +++ b/hpgsql-benchmarks/src/Main.hs @@ -154,7 +154,13 @@ main = do putStrLn "IMPORTANT: all measurements collected over 10 runs of each benchmark" when (numConcurrentConnections > 1) $ putStrLn $ "IMPORTANT: all benchmarks except COPY involve running the benchmarked query in " ++ show numConcurrentConnections ++ " connections in parallel" + + -- Warm up postgres with a generate_series query and GC before tests + warmupConn <- hpgsqlConnect + void $ Hpgsql.execute warmupConn "SELECT * FROM generate_series(1,100000)" + Hpgsql.Connection.closeGracefully warmupConn performBlockingMajorGC + statsBefore <- getRTSStats hspecWith defaultConfig {configFormat = Just (formatterToFormat silent)} $ do describe "Parsing 13-column rows into a List" $ do diff --git a/hpgsql-tests/CopySpec.hs b/hpgsql-tests/CopySpec.hs index 15ccae7..c58ca00 100644 --- a/hpgsql-tests/CopySpec.hs +++ b/hpgsql-tests/CopySpec.hs @@ -2,7 +2,7 @@ module CopySpec where import Control.Monad (forM_) import Control.Monad.IO.Class (liftIO) -import Data.Int (Int32) +import Data.Int (Int32, Int64) import Data.Text (Text) import qualified Data.Text as Text import qualified Data.Text.Encoding as TE @@ -38,36 +38,37 @@ spec = do "putCopyError" copyError -genRows :: Gen.Gen [(Int32, Text)] +genRows :: Gen.Gen [(Int32, Text, Int64)] genRows = do numRows <- Gen.int (Gen.linear 0 1000) names <- Gen.list (Gen.singleton numRows) $ Gen.text (Gen.linear 1 50) Gen.alphaNum - pure $ zip [1 ..] names + numbers <- Gen.list (Gen.singleton numRows) $ Gen.int64 (Gen.linear (-100) 100) + pure $ zip3 [1 ..] names numbers copyTextFmtStatementSucceeding :: HPgConnection -> PropertyT IO () copyTextFmtStatementSucceeding conn = hedgehog $ do rows <- Gen.forAll genRows result <- liftIO $ withRollback conn $ do - execute_ conn "CREATE UNLOGGED TABLE copy_test0 (id INT NOT NULL, name TEXT NOT NULL)" + execute_ conn "CREATE UNLOGGED TABLE copy_test0 (id INT NOT NULL, name TEXT NOT NULL, some_num BIGINT)" withCopy_ conn "COPY copy_test0 FROM STDIN WITH (FORMAT CSV);" - ( forM_ rows $ \(eid, ename) -> - putCopyData conn $ TE.encodeUtf8 $ Text.pack (show eid) <> "," <> ename <> "\n" + ( forM_ rows $ \(eid, ename, somenum) -> + putCopyData conn $ TE.encodeUtf8 $ Text.pack (show eid) <> "," <> ename <> "," <> Text.pack (show somenum) <> "\n" ) - query conn "SELECT id, name FROM copy_test0 ORDER BY id" + query conn "SELECT id, name, some_num FROM copy_test0 ORDER BY id" result === rows copyBinaryFmtStatementSucceeding :: HPgConnection -> PropertyT IO () copyBinaryFmtStatementSucceeding conn = hedgehog $ do rows <- Gen.forAll genRows result <- liftIO $ withRollback conn $ do - execute_ conn "CREATE UNLOGGED TABLE copy_test1 (id INT NOT NULL, name TEXT NOT NULL)" + execute_ conn "CREATE UNLOGGED TABLE copy_test1 (id INT NOT NULL, name TEXT NOT NULL, some_num BIGINT)" copyFrom conn "COPY copy_test1 FROM STDIN WITH (FORMAT BINARY);" rows - query conn "SELECT id, name FROM copy_test1 ORDER BY id" + query conn "SELECT id, name, some_num FROM copy_test1 ORDER BY id" result === rows copyError :: HPgConnection -> IO () diff --git a/hpgsql-tests/EncodingDecodingSpec.hs b/hpgsql-tests/EncodingDecodingSpec.hs index ca24dda..37e63cd 100644 --- a/hpgsql-tests/EncodingDecodingSpec.hs +++ b/hpgsql-tests/EncodingDecodingSpec.hs @@ -31,6 +31,7 @@ import qualified Data.Vector as Vector import DbUtils ( aroundConn, irrecoverableErrorWithMsgAndStmt, + testConnInfo, withRollback, ) import GHC.Float (float2Double) @@ -40,7 +41,7 @@ import qualified Hedgehog as Gen import qualified Hedgehog.Gen as Gen import qualified Hedgehog.Range as Gen import Hpgsql -import Hpgsql.Connection (refreshTypeInfoCache) +import Hpgsql.Connection (ConnectOpts (..), connect, connectOpts, defaultConnectOpts, refreshTypeInfoCache, withConnectionOpts) import Hpgsql.Encoding (EncodingContext (..), FieldDecoder (..), FieldEncoder (..), FieldInfo (..), FromPgField (..), FromPgRow (..), LowerCasedPgEnum (..), RowEncoder (..), ToPgField (..), ToPgRow (..), compositeTypeDecoder, compositeTypeEncoder, nullableField, rawBytesFieldDecoder, singleField, typeFieldDecoder, typeFieldEncoder, typeMustBeNamed, typeOidWithName) import Hpgsql.Pipeline (pipeline, pipelineWith, runPipeline) import Hpgsql.Query (mkQuery, sql, vALUES) @@ -156,16 +157,22 @@ spec = parallel $ do it "Generically derived types round-trip" queryGenericallyDerivedTypesRoundTrip + it + "0-columns results can be decoded" + zeroColumnsResults + +zeroColumnsResults :: IO () +zeroColumnsResults = do + hpgsqlConnInfo <- testConnInfo + -- This test is important to test the "slow" decoding path of `decodeDataRow` + -- in BinarySerializer.hs. The number of rows needs to be a bit large + -- and the recvChunkSize pretty small for that code path to be exercised, + -- as per some debug printing. + withConnectionOpts defaultConnectOpts {recvChunkSize = 5} hpgsqlConnInfo 10 $ \conn -> do + execute conn "SELECT FROM generate_series(1,601)" `shouldReturn` 601 valuesRoundTrip :: HPgConnection -> IO () valuesRoundTrip conn = do - -- TODO: Property-based test to generate the values - -- TODO: Include NULLs - -- TODO: Test +-infinity for types where we can - -- TODO: Test all types in the regions of values close to `minBound`, 0, and `maxBound` - -- TODO: Test floats, timestamptz and other very granular but discrete type in the regions of values - -- close to `minBound`, 0, and `maxBound`, with e.g. microsecond precision/fractional values - -- TODO: Test +-Infinity and NaN for floats and doubles let row = ((-49) :: Int, False :: Bool, 2 :: Int16, 3 :: Int32, fromGregorian 1900 02 28, 42 :: Int64, UTCTime (fromGregorian 1999 12 31) 0, '意' :: Char, '&' :: Char, CalendarDiffTime 3 86403, Aeson.Null) queryWith rowDecoder conn (mkQuery "SELECT $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11" row) `shouldReturn` [row] diff --git a/hpgsql/hpgsql.cabal b/hpgsql/hpgsql.cabal index ec89502..d20fb9c 100644 --- a/hpgsql/hpgsql.cabal +++ b/hpgsql/hpgsql.cabal @@ -43,6 +43,7 @@ library Hpgsql.Types other-modules: Hpgsql.Base + Hpgsql.Encoding.BinarySerializer Hpgsql.Internal Hpgsql.LanguageHaskell.FromThExtension Hpgsql.LanguageHaskell.GhcParserOpts @@ -102,7 +103,6 @@ library base >= 4.18 && < 4.22, bytestring >= 0.11 && < 0.13, case-insensitive >= 1.2 && < 1.3, - cereal >= 0.5 && < 0.6, containers >= 0.6 && < 0.8, crypton >= 1.0.0 && < 1.1, memory >= 0.18.0 && < 0.19, diff --git a/hpgsql/src/Hpgsql/Builder.hs b/hpgsql/src/Hpgsql/Builder.hs index eeb8e20..93396de 100644 --- a/hpgsql/src/Hpgsql/Builder.hs +++ b/hpgsql/src/Hpgsql/Builder.hs @@ -1,12 +1,10 @@ -module Hpgsql.Builder where - --- \| This module replicates parts of the API of Data.ByteString.Builder but its own +-- | This module replicates parts of the API of Data.ByteString.Builder but its own -- builder is length-aware, which makes other parts of the code a little bit nicer. -- In COPY benchmarks, this module was introduced in a commit (together with other -- changes, like replacing `Maybe` with `BinaryField` in `ToPgField`) that barely -- changed memory usage and runtime. -- The benefits are exclusively for code readability, then. --- \| +module Hpgsql.Builder where import Data.ByteString (ByteString) import qualified Data.ByteString as BS @@ -23,7 +21,9 @@ instance Show BinaryField where show SqlNull = "NULL" show (NotNull bs) = show bs -data LengthAwareBuilder = LengthAwareBuilder !Int32 !Builder.Builder +-- | The lazy (instead of strict/with a bang) Builder (second arg) makes +-- our copyFromS benchmark run ~4.3% faster and allocate ~3.8% less total memory. +data LengthAwareBuilder = LengthAwareBuilder !Int32 Builder.Builder type Builder = LengthAwareBuilder diff --git a/hpgsql/src/Hpgsql/Connection.hs b/hpgsql/src/Hpgsql/Connection.hs index a831119..84795c4 100644 --- a/hpgsql/src/Hpgsql/Connection.hs +++ b/hpgsql/src/Hpgsql/Connection.hs @@ -8,6 +8,7 @@ module Hpgsql.Connection closeForcefully, connectionIsClosed, ConnectionString (..), + ConnectOpts (..), parseLibpqConnectionString, ResetConnectionOpts (..), resetConnectionState, @@ -48,7 +49,7 @@ import Data.Text (Text) import qualified Data.Text as Text import Data.Text.Encoding (encodeUtf8) import Hpgsql.Internal (closeForcefully, closeGracefully, connect, connectOpts, connectionIsClosed, defaultConnectOpts, getBackendPid, getParameterStatus, refreshTypeInfoCache, resetConnectionState, resetTypeInfoCache, withConnection, withConnectionOpts) -import Hpgsql.InternalTypes (ConnectionString (..), ResetConnectionOpts (..)) +import Hpgsql.InternalTypes (ConnectOpts (..), ConnectionString (..), ResetConnectionOpts (..)) import Network.URI ( URI (..), URIAuth (..), diff --git a/hpgsql/src/Hpgsql/Encoding.hs b/hpgsql/src/Hpgsql/Encoding.hs index 3636db3..b03a83c 100644 --- a/hpgsql/src/Hpgsql/Encoding.hs +++ b/hpgsql/src/Hpgsql/Encoding.hs @@ -87,7 +87,6 @@ import Data.Monoid (Sum (..)) import Data.Proxy (Proxy (..)) import Data.Ratio (Ratio) import Data.Scientific (Scientific (..), floatingOrInteger, scientific) -import qualified Data.Serialize as Cereal import Data.Text (Text) import qualified Data.Text as Text import Data.Text.Encoding (decodeUtf8, encodeUtf8) @@ -100,13 +99,13 @@ import Data.UUID.Types (UUID) import qualified Data.UUID.Types as UUID import Data.Vector (Vector) import qualified Data.Vector as Vector -import Data.Word (Word32, Word64) -import GHC.Float (castDoubleToWord64, castFloatToWord32, castWord32ToFloat, castWord64ToDouble, expt, float2Double) +import GHC.Float (castWord32ToFloat, castWord64ToDouble, expt, float2Double) import GHC.Generics (C, D, Generic (..), K1 (..), M1 (..), Meta (MetaCons), U1 (..), (:*:) (..), (:+:) (..)) import GHC.TypeLits (KnownSymbol, TypeError, symbolVal) import qualified GHC.TypeLits as TypeLits import Hpgsql.Builder (BinaryField (..)) import qualified Hpgsql.Builder as Builder +import qualified Hpgsql.Encoding.BinarySerializer as BinSer import qualified Hpgsql.SimpleParser as Parser import Hpgsql.Time (Unbounded (..)) import Hpgsql.TypeInfo (EncodingContext (..), Oid (..), TypeDetails (..), TypeInfo (..), boolOid, byteaOid, charOid, dateOid, float4Oid, float8Oid, int2Oid, int4Oid, int8Oid, intervalOid, jsonOid, jsonbOid, lookupTypeByName, lookupTypeByOid, nameOid, numericOid, oidOid, textOid, timeOid, timestampOid, timestamptzOid, uuidOid, varcharOid, voidOid) @@ -163,7 +162,7 @@ singleField (FieldDecoder {..}) = [singleColInfo] -> let decode = fieldValueDecoder singleColInfo in do - lenNextCol <- fromIntegral <$> int32Parser + lenNextCol <- fromIntegral <$> Parser.takeInt32BE nextColBs <- if lenNextCol >= 0 then @@ -179,9 +178,6 @@ singleField (FieldDecoder {..}) = numExpectedColumns = 1 } -int32Parser :: Parser.Parser Int32 -int32Parser = either fail pure . Cereal.decode @Int32 =<< Parser.take 4 - class FromPgField a where fieldDecoder :: FieldDecoder a @@ -216,12 +212,12 @@ compositeTypeDecoder (RowDecoder {..}) = parserForRecord encodingContext = do -- From https://github.com/postgres/postgres/blob/50ba65e73325cf55fedb3e1f14673d816726923b/src/backend/utils/adt/rowtypes.c#L687 -- we can see a composite type's binary representation consists of: number of columns (Int32) + for_each_column { OID (Int32) + size_or_minus_1 (Int32) + Bytes } - numCols <- fromIntegral <$> int32Parser + numCols <- fromIntegral <$> Parser.takeInt32BE unless (numCols == numExpectedColumns) $ fail $ "Composite type has " ++ show numCols ++ " attributes but parser expected " ++ show numExpectedColumns let mkColInfo oid = FieldInfo oid Nothing encodingContext cols <- replicateM numCols $ do - !oid <- Oid . fromIntegral <$> int32Parser - (sizeBs, !size) <- Parser.match $ fromIntegral <$> int32Parser + !oid <- Oid . fromIntegral <$> Parser.takeInt32BE + (sizeBs, !size) <- Parser.match $ fromIntegral <$> Parser.takeInt32BE !bs <- Parser.take (max 0 size) pure (oid, sizeBs <> bs) let typecheckedCols = rowColumnsTypeCheck (map (mkColInfo . fst) cols) @@ -337,21 +333,21 @@ instance ToPgField Int16 where fieldEncoder = FieldEncoder { toTypeOid = \_ -> Just int2Oid, - toPgField = \_ -> \n -> NotNull $ Cereal.encode n + toPgField = \_ -> \n -> NotNull $ BinSer.encodeInt16BE n } instance ToPgField Int32 where fieldEncoder = FieldEncoder { toTypeOid = \_ -> Just int4Oid, - toPgField = \_ -> \n -> NotNull $ Cereal.encode n + toPgField = \_ -> \n -> NotNull $ BinSer.encodeInt32BE n } instance ToPgField Int64 where fieldEncoder = FieldEncoder { toTypeOid = \_ -> Just int8Oid, - toPgField = \_ -> \n -> NotNull $ Cereal.encode n + toPgField = \_ -> \n -> NotNull $ BinSer.encodeInt64BE n } instance ToPgField Integer where @@ -374,7 +370,7 @@ instance ToPgField Oid where fieldEncoder = FieldEncoder { toTypeOid = \_ -> Just oidOid, - toPgField = \_ -> \n -> NotNull $ Cereal.encode @Int32 $ fromIntegral n + toPgField = \_ -> \n -> NotNull $ BinSer.encodeInt32BE $ fromIntegral n } instance ToPgField Scientific where @@ -382,7 +378,7 @@ instance ToPgField Scientific where FieldEncoder { toTypeOid = \_ -> Just numericOid, toPgField = \_ -> \n -> - let sign = Cereal.encode @Int16 $ if n >= 0 then 0 else 0x4000 + let sign = BinSer.encodeInt16BE $ if n >= 0 then 0 else 0x4000 -- The number is coeff * 10^exp, but we want it in base-10000 so we convert it to -- new_coeff * 10^new_exp with new_exp a multiple of 4 base10000Expon = 4 * (base10Exponent n `div` 4) @@ -390,8 +386,8 @@ instance ToPgField Scientific where ndigits, weight :: Int16 digits :: ByteString (ndigits, weight, digits) = calculateDigits 0 0 (abs base10000Coeff) "" - dscale = Cereal.encode @Int16 (abs $ fromIntegral base10000Expon) -- More than necessary, but safe? - in NotNull $ Cereal.encode ndigits <> Cereal.encode (weight - 1 + fromIntegral (base10000Expon `div` 4)) <> sign <> dscale <> digits + dscale = BinSer.encodeInt16BE (abs $ fromIntegral base10000Expon) -- More than necessary, but safe? + in NotNull $ BinSer.encodeInt16BE ndigits <> BinSer.encodeInt16BE (weight - 1 + fromIntegral (base10000Expon `div` 4)) <> sign <> dscale <> digits } where calculateDigits :: Int16 -> Int16 -> Integer -> BS.ByteString -> (Int16, Int16, BS.ByteString) @@ -402,28 +398,27 @@ instance ToPgField Scientific where (ndigitsSoFar + 1) (weightSoFar + 1) quotient - (Cereal.encode @Int16 rest <> encodedDigits) + (BinSer.encodeInt16BE rest <> encodedDigits) instance ToPgField Float where fieldEncoder = FieldEncoder { toTypeOid = \_ -> Just float4Oid, - toPgField = \_ -> \n -> NotNull $ Cereal.encode @Word32 $ castFloatToWord32 n + toPgField = \_ -> \n -> NotNull $ BinSer.encodeFloat n } instance ToPgField Double where fieldEncoder = FieldEncoder { toTypeOid = \_ -> Just float8Oid, - toPgField = \_ -> \n -> NotNull $ Cereal.encode @Word64 $ castDoubleToWord64 n + toPgField = \_ -> \n -> NotNull $ BinSer.encodeDouble n } instance ToPgField Bool where - -- TODO: Cereal.encode seems to work, but reference the documentation that shows how bools are encoded fieldEncoder = FieldEncoder { toTypeOid = \_ -> Just boolOid, - toPgField = \_ n -> NotNull $ Cereal.encode @Bool $ n + toPgField = \_ n -> NotNull $ BinSer.encodePgBoolean n } instance ToPgField Day where @@ -433,7 +428,7 @@ instance ToPgField Day where FieldEncoder { toTypeOid = \_ -> Just dateOid, -- TODO: Catch integer overflow and do what? - toPgField = \_ d -> NotNull $ Cereal.encode @Int32 $ fromIntegral $ diffDays d (fromGregorian 2000 1 1) + toPgField = \_ d -> NotNull $ BinSer.encodeInt32BE $ fromIntegral $ diffDays d (fromGregorian 2000 1 1) } instance ToPgField (Unbounded Day) where @@ -442,9 +437,9 @@ instance ToPgField (Unbounded Day) where in FieldEncoder { toTypeOid = fe.toTypeOid, toPgField = \encCtx -> \case - NegInfinity -> NotNull $ Cereal.encode @Int32 minBound + NegInfinity -> NotNull $ BinSer.encodeInt32BE minBound Finite v -> fe.toPgField encCtx v - PosInfinity -> NotNull $ Cereal.encode @Int32 maxBound + PosInfinity -> NotNull $ BinSer.encodeInt32BE maxBound } instance ToPgField CalendarDiffTime where @@ -453,7 +448,7 @@ instance ToPgField CalendarDiffTime where { toTypeOid = \_ -> Just intervalOid, toPgField = \_ CalendarDiffTime {..} -> let (days :: Int32, timeUnderOneDay) = ctTime `divMod'` 86_400 - in NotNull $ Cereal.encode @(Int64, Int32, Int32) (round $ timeUnderOneDay * 1_000_000, days, fromIntegral ctMonths) + in NotNull $ BinSer.encodeInt64BE (round $ timeUnderOneDay * 1_000_000) <> BinSer.encodeInt32BE days <> BinSer.encodeInt32BE (fromIntegral ctMonths) } instance ToPgField NominalDiffTime where @@ -461,7 +456,7 @@ instance ToPgField NominalDiffTime where FieldEncoder { toTypeOid = \_ -> Just intervalOid, toPgField = \_ ndt -> - NotNull $ Cereal.encode @(Int64, Int32, Int32) (round $ ndt * 1_000_000, 0, 0) + NotNull $ BinSer.encodeInt64BE (round $ ndt * 1_000_000) <> BinSer.encodeInt32BE 0 <> BinSer.encodeInt32BE 0 } instance ToPgField UTCTime where @@ -472,7 +467,7 @@ instance ToPgField UTCTime where toPgField = \_ (UTCTime parsedDate timeinday) -> let day :: Int64 = fromInteger $ parsedDate `diffDays` fromJulian 1999 12 19 totalusecs :: Int64 = 86_400_000_000 * day + fromInteger (diffTimeToPicoseconds timeinday `div` 1_000_000) - in NotNull $ Cereal.encode @Int64 totalusecs + in NotNull $ BinSer.encodeInt64BE totalusecs } instance ToPgField (Unbounded UTCTime) where @@ -481,9 +476,9 @@ instance ToPgField (Unbounded UTCTime) where in FieldEncoder { toTypeOid = fe.toTypeOid, toPgField = \encCtx -> \case - NegInfinity -> NotNull $ Cereal.encode @Int64 minBound + NegInfinity -> NotNull $ BinSer.encodeInt64BE minBound Finite v -> fe.toPgField encCtx v - PosInfinity -> NotNull $ Cereal.encode @Int64 maxBound + PosInfinity -> NotNull $ BinSer.encodeInt64BE maxBound } instance ToPgField ZonedTime where @@ -500,9 +495,9 @@ instance ToPgField (Unbounded ZonedTime) where in FieldEncoder { toTypeOid = fe.toTypeOid, toPgField = \encCtx -> \case - NegInfinity -> NotNull $ Cereal.encode @Int64 minBound + NegInfinity -> NotNull $ BinSer.encodeInt64BE minBound Finite v -> fe.toPgField encCtx v - PosInfinity -> NotNull $ Cereal.encode @Int64 maxBound + PosInfinity -> NotNull $ BinSer.encodeInt64BE maxBound } instance ToPgField LocalTime where @@ -512,7 +507,7 @@ instance ToPgField LocalTime where toPgField = \_ (LocalTime localDay localTimeOfDay) -> let day :: Int64 = fromInteger $ localDay `diffDays` fromJulian 1999 12 19 totalusecs :: Int64 = 86_400_000_000 * day + fromInteger (diffTimeToPicoseconds (timeOfDayToTime localTimeOfDay) `div` 1_000_000) - in NotNull $ Cereal.encode @Int64 totalusecs + in NotNull $ BinSer.encodeInt64BE totalusecs } instance ToPgField TimeOfDay where @@ -521,7 +516,7 @@ instance ToPgField TimeOfDay where { toTypeOid = \_ -> Just timeOid, toPgField = \_ tod -> let usecs :: Int64 = fromInteger $ diffTimeToPicoseconds (timeOfDayToTime tod) `div` 1_000_000 - in NotNull $ Cereal.encode @Int64 usecs + in NotNull $ BinSer.encodeInt64BE usecs } instance ToPgField Char where @@ -684,15 +679,6 @@ instance (ToPgField a, ToPgField b, ToPgField c) => ToPgRow (a, b, c) where instance (ToPgField a, ToPgField b, ToPgField c, ToPgField d) => ToPgRow (a, b, c, d) where rowEncoder = divide (\(a, b, c, d) -> ((a, b), (c, d))) rowEncoder rowEncoder --- This instance implements toBinaryCopyBytes as well because we did this --- to test if this method can help improve performance of COPY in our --- benchmarks. We found that it can, but we didn't bother yet implementing --- this for other types. --- toBinaryCopyBytes encCtx = \(a, b, c, d) -> Builder.int16BE 4 <> toPgFieldWithSize a <> toPgFieldWithSize b <> toPgFieldWithSize c <> toPgFieldWithSize d --- where --- toPgFieldWithSize :: (ToPgField x) => x -> Builder.Builder --- toPgFieldWithSize v = Builder.binaryField $ toPgField encCtx v - instance (ToPgField a, ToPgField b, ToPgField c, ToPgField d, ToPgField e) => ToPgRow (a, b, c, d, e) where rowEncoder = divide (\(a, b, c, d, e) -> ((a, b, c), (d, e))) rowEncoder rowEncoder @@ -733,9 +719,9 @@ haskellIntOids :: [Oid] -- | Big-Endian binary encoder for Haskell's `Data.Int`, which is machine-dependent. binaryIntEncoder :: Int -> BinaryField binaryIntEncoder - | haskellIntOid == int8Oid = NotNull . Cereal.encode @Int64 . fromIntegral - | haskellIntOid == int4Oid = NotNull . Cereal.encode @Int32 . fromIntegral - | otherwise = NotNull . Cereal.encode @Int16 . fromIntegral + | haskellIntOid == int8Oid = NotNull . BinSer.encodeInt64BE . fromIntegral + | haskellIntOid == int4Oid = NotNull . BinSer.encodeInt32BE . fromIntegral + | otherwise = NotNull . BinSer.encodeInt16BE . fromIntegral -- | Big-Endian binary decoder for Haskell's various IntXX types. binaryIntDecoder :: forall a. (Integral a, Bounded a) => Oid -> ByteString -> Either String a @@ -747,17 +733,17 @@ binaryIntDecoder typOid = \bs -> maxBoundPgType :: Integer intDecoder :: ByteString -> Either String a (maxBoundPgType, intDecoder) - | typOid == int8Oid = (fromIntegral $ maxBound @Int64, fmap fromIntegral . Cereal.decode @Int64) - | typOid == int4Oid = (fromIntegral $ maxBound @Int32, fmap fromIntegral . Cereal.decode @Int32) - | typOid == int2Oid = (fromIntegral $ maxBound @Int16, fmap fromIntegral . Cereal.decode @Int16) + | typOid == int8Oid = (fromIntegral $ maxBound @Int64, fmap fromIntegral . BinSer.decodeInt64BE) + | typOid == int4Oid = (fromIntegral $ maxBound @Int32, fmap fromIntegral . BinSer.decodeInt32BE) + | typOid == int2Oid = (fromIntegral $ maxBound @Int16, fmap fromIntegral . BinSer.decodeInt16BE) | otherwise = error "Bug in Hpgsql. Decoding binary integral type not an int2, int4 or int8" doesFit = maxBoundPgType <= fromIntegral (maxBound @a) binaryFloat4Decoder :: ByteString -> Float -binaryFloat4Decoder = castWord32ToFloat . either error id . Cereal.decode @Word32 +binaryFloat4Decoder = castWord32ToFloat . either error id . BinSer.decodeWord32BE binaryFloat8Decoder :: ByteString -> Double -binaryFloat8Decoder = castWord64ToDouble . either error id . Cereal.decode @Word64 +binaryFloat8Decoder = castWord64ToDouble . either error id . BinSer.decodeWord64BE parsePgType :: [Oid] -> (Maybe ByteString -> Either String a) -> FieldDecoder a parsePgType !requiredTypeOids !fieldValueDecoder = @@ -890,11 +876,11 @@ typeMustBeNamed typName = \fieldInfo -> scientificDecoder :: Bool -> Parser.Parser Scientific scientificDecoder mustBeInteger = do - ndigits <- int16Parser - weight <- int16Parser - sign <- int16Parser -- 0x0000 is positive, 0x4000 is negative, 0xC000 is NAN, 0xD000 is Positive Infinity, 0xF000 is Negative Infinity + ndigits <- Parser.takeInt16BE + weight <- Parser.takeInt16BE + sign <- Parser.takeInt16BE -- 0x0000 is positive, 0x4000 is negative, 0xC000 is NAN, 0xD000 is Positive Infinity, 0xF000 is Negative Infinity unless (sign == 0x0000 || sign == 0x4000) $ fail "NaN, positive or negative infinities cannot be decoded into Integer or Scientific" - !dscale <- int16Parser + !dscale <- Parser.takeInt16BE when (mustBeInteger && dscale /= 0) $ fail "Decoding into `Integer` requires explicit casting with `numeric(X,0)` to force integral values" valueAbs <- parseAndMult ndigits (fromIntegral weight * 4) 0 pure $ (if sign == 0x0000 then 1 else (-1)) * valueAbs @@ -902,7 +888,7 @@ scientificDecoder mustBeInteger = do parseAndMult :: Int16 -> Int -> Scientific -> Parser.Parser Scientific parseAndMult 0 _ !val = pure val parseAndMult !ndigitsLeft !currexpon !val = do - !digit <- fromIntegral <$> int16Parser + !digit <- fromIntegral <$> Parser.takeInt16BE parseAndMult (ndigitsLeft - 1) (currexpon - 4) (val + scientific digit currexpon) instance FromPgField Scientific where @@ -928,7 +914,7 @@ instance FromPgField (Ratio Integer) where fieldDecoder = toRational <$> fieldDecoder @Scientific binaryTrue :: ByteString -binaryTrue = Cereal.encode True +binaryTrue = BinSer.encodePgBoolean True instance FromPgField Bool where fieldDecoder = parsePgType [boolOid] $ \case @@ -1003,7 +989,7 @@ instance FromPgField UTCTime where fieldDecoder = parsePgType [timestamptzOid] $ \case Just bs -> do -- See https://github.com/postgres/postgres/blob/50cb7505b3010736b9a7922e903931534785f3aa/src/backend/utils/adt/timestamp.c#L1909 - totalusecs <- Cereal.decode @Int64 bs + totalusecs <- BinSer.decodeInt64BE bs let (day, timeusecs) = totalusecs `divMod` 86_400_000_000 -- USECS per day parsedDate = addJulianDurationClip (CalendarDiffDays 0 (fromIntegral day)) $ fromJulian 1999 12 19 Right $ UTCTime parsedDate (picosecondsToDiffTime $ fromIntegral timeusecs * 1_000_000) @@ -1013,7 +999,7 @@ instance FromPgField (Unbounded UTCTime) where fieldDecoder = parsePgType [timestamptzOid] $ \case Just bs -> do -- See https://github.com/postgres/postgres/blob/50cb7505b3010736b9a7922e903931534785f3aa/src/backend/utils/adt/timestamp.c#L1909 - totalusecs <- Cereal.decode @Int64 bs + totalusecs <- BinSer.decodeInt64BE bs Right $ if totalusecs == minBound then NegInfinity @@ -1030,7 +1016,7 @@ instance FromPgField ZonedTime where fieldDecoder = parsePgType [timestamptzOid] $ \case Just bs -> do -- See https://github.com/postgres/postgres/blob/50cb7505b3010736b9a7922e903931534785f3aa/src/backend/utils/adt/timestamp.c#L1909 - totalusecs <- Cereal.decode @Int64 bs + totalusecs <- BinSer.decodeInt64BE bs let (day, timeusecs) = totalusecs `divMod` 86_400_000_000 -- USECS per day parsedDate = addJulianDurationClip (CalendarDiffDays 0 (fromIntegral day)) $ fromJulian 1999 12 19 Right $ utcToZonedTime utc $ UTCTime parsedDate (picosecondsToDiffTime $ fromIntegral timeusecs * 1_000_000) @@ -1040,7 +1026,7 @@ instance FromPgField (Unbounded ZonedTime) where fieldDecoder = parsePgType [timestamptzOid] $ \case Just bs -> do -- See https://github.com/postgres/postgres/blob/50cb7505b3010736b9a7922e903931534785f3aa/src/backend/utils/adt/timestamp.c#L1909 - totalusecs <- Cereal.decode @Int64 bs + totalusecs <- BinSer.decodeInt64BE bs Right $ if totalusecs == minBound then NegInfinity @@ -1056,7 +1042,7 @@ instance FromPgField (Unbounded ZonedTime) where instance FromPgField LocalTime where fieldDecoder = parsePgType [timestampOid] $ \case Just bs -> do - totalusecs <- Cereal.decode @Int64 bs + totalusecs <- BinSer.decodeInt64BE bs let (day, timeusecs) = totalusecs `divMod` 86_400_000_000 -- USECS per day parsedDate = addJulianDurationClip (CalendarDiffDays 0 (fromIntegral day)) $ fromJulian 1999 12 19 Right $ LocalTime parsedDate (timeToTimeOfDay $ picosecondsToDiffTime $ fromIntegral timeusecs * 1_000_000) @@ -1065,7 +1051,7 @@ instance FromPgField LocalTime where instance FromPgField TimeOfDay where fieldDecoder = parsePgType [timeOid] $ \case Just bs -> do - usecs <- Cereal.decode @Int64 bs + usecs <- BinSer.decodeInt64BE bs Right $ timeToTimeOfDay $ picosecondsToDiffTime $ fromIntegral usecs * 1_000_000 Nothing -> Left "Cannot decode SQL null as the Haskell TimeOfDay type. Use a `Maybe TimeOfDay`" @@ -1075,7 +1061,7 @@ instance FromPgField Day where -- There is a very specific conversion function for these, which I poorly translated to Haskell -- https://github.com/postgres/postgres/blob/799959dc7cf0e2462601bea8d07b6edec3fa0c4f/src/backend/utils/adt/datetime.c#L321 -- But I found a simpler way to do this. Let's see if it works in our property based tests - jd <- Cereal.decode @Int32 bs + jd <- BinSer.decodeInt32BE bs Right $ addJulianDurationClip (CalendarDiffDays 0 (fromIntegral jd - 13)) $ fromJulian 2000 01 01 Nothing -> Left "Cannot decode SQL null as the Haskell Day type. Use a `Maybe Day`" @@ -1085,7 +1071,7 @@ instance FromPgField (Unbounded Day) where -- There is a very specific conversion function for these, which I poorly translated to Haskell -- https://github.com/postgres/postgres/blob/799959dc7cf0e2462601bea8d07b6edec3fa0c4f/src/backend/utils/adt/datetime.c#L321 -- But I found a simpler way to do this. Let's see if it works in our property based tests - jd <- Cereal.decode @Int32 bs + jd <- BinSer.decodeInt32BE bs Right $ if jd == minBound then NegInfinity @@ -1099,7 +1085,9 @@ instance FromPgField (Unbounded Day) where instance FromPgField CalendarDiffTime where fieldDecoder = parsePgType [intervalOid] $ \case Just bs -> do - (nMicrosecs :: Int64, nDays :: Int32, nMonths :: Int32) <- Cereal.decode bs + nMicrosecs <- BinSer.decodeInt64BE bs + nDays <- BinSer.decodeInt32BE (BS.drop 8 bs) + nMonths <- BinSer.decodeInt32BE (BS.drop 12 bs) Right $ CalendarDiffTime {ctMonths = fromIntegral nMonths, ctTime = secondsToNominalDiffTime (fromIntegral nDays * 86400) + realToFrac (picosecondsToDiffTime (fromIntegral nMicrosecs * 1_000_000))} Nothing -> Left "Cannot decode SQL null as the Haskell CalendarDiffTime type. Use a `Maybe CalendarDiffTime`" @@ -1171,33 +1159,30 @@ instance {-# OVERLAPPING #-} forall a. (FromPgField a) => FromPgField (Vector (V !elementParser = fieldDecoder @a arrayParser :: EncodingContext -> Parser.Parser (Vector (Vector a)) arrayParser encodingContext = do - !ndim <- int32Parser - !_hasNull <- int32Parser - !elementTypeOid :: Oid <- Oid . fromIntegral <$> int32Parser + !ndim <- Parser.takeInt32BE + !_hasNull <- Parser.takeInt32BE + !elementTypeOid :: Oid <- Oid . fromIntegral <$> Parser.takeInt32BE let !elementColInfo = FieldInfo elementTypeOid Nothing encodingContext when (ndim /= 2) $ fail $ "TODO: No support for " ++ show ndim ++ "-dimensional arrays in Hpgsql. Got array with ndim=" ++ show ndim unless (elementParser.allowedPgTypes elementColInfo) $ fail $ "Array contains elements of type OID " ++ show elementTypeOid ++ " but decoder does not handle that type" numRows <- do - !dim_i :: Int <- fromIntegral <$> int32Parser - !_lb_i <- int32Parser + !dim_i :: Int <- fromIntegral <$> Parser.takeInt32BE + !_lb_i <- Parser.takeInt32BE pure dim_i lengthEachRow <- do - !dim_i :: Int <- fromIntegral <$> int32Parser - !_lb_i <- int32Parser + !dim_i :: Int <- fromIntegral <$> Parser.takeInt32BE + !_lb_i <- Parser.takeInt32BE pure dim_i Vector.replicateM numRows $ do Vector.replicateM lengthEachRow $ do - size :: Int <- fromIntegral <$> int32Parser + size :: Int <- fromIntegral <$> Parser.takeInt32BE elementBs <- if size == (-1) then pure Nothing else Just <$> Parser.take size case elementParser.fieldValueDecoder elementColInfo elementBs of Left err -> fail $ "Error parsing array element: " ++ show err Right el -> pure el -int16Parser :: Parser.Parser Int16 -int16Parser = either fail pure . Cereal.decode @Int16 =<< Parser.take 2 - -- | Derives `FromPgRow` generically. genericFromPgRow :: forall a. (Generic a, ProductTypeDecoder (Rep a)) => RowDecoder a genericFromPgRow = to <$> genRowDecoder @(Rep a) @@ -1333,14 +1318,14 @@ toPgVectorField encCtx = encodeElement el = Builder.binaryField $ fe.toPgField encCtx el Oid elemOid = fromMaybe (Oid 0) (fe.toTypeOid encCtx) in \vec -> - let ndim = Builder.byteString $ Cereal.encode @Int32 1 + let ndim = Builder.int32BE 1 -- Postgres seems to build the "has_nulls" flag itself in the ReadArrayBinary function at https://github.com/postgres/postgres/blob/aa7f9493a02f5981c09b924323f0e7a58a32f2ed/src/backend/utils/adt/arrayfuncs.c#L1429, so we can just set it to 0 - hasNull = Builder.byteString $ Cereal.encode @Int32 0 - -- hasNull = Builder.byteString $ Cereal.encode @Int32 (if Vector.any (\e -> toPgField e == Nothing) vec then 1 else 0) - elemOidBs = Builder.byteString $ Cereal.encode @Int32 elemOid - lb1 = Builder.byteString $ Cereal.encode @Int32 1 + hasNull = Builder.byteString $ BinSer.encodeInt32BE 0 + -- hasNull = Builder.byteString $ BinSer.encodeInt32BE (if Vector.any (\e -> toPgField e == Nothing) vec then 1 else 0) + elemOidBs = Builder.byteString $ BinSer.encodeInt32BE elemOid + lb1 = Builder.byteString $ BinSer.encodeInt32BE 1 (Sum len, encodedElements) = foldMap (\el -> (Sum 1, encodeElement el)) vec - dim1 = Builder.byteString $ Cereal.encode @Int32 len + dim1 = Builder.byteString $ BinSer.encodeInt32BE len fullBs = ndim <> hasNull <> elemOidBs <> dim1 <> lb1 <> encodedElements in NotNull (Builder.toStrictByteString fullBs) @@ -1361,19 +1346,19 @@ arrayField !replicateFunction !elementParser = where arrayParser :: EncodingContext -> Parser.Parser (f a) arrayParser encodingContext = do - !ndim <- int32Parser - !_hasNull <- int32Parser - !elementTypeOid :: Oid <- Oid . fromIntegral <$> int32Parser + !ndim <- Parser.takeInt32BE + !_hasNull <- Parser.takeInt32BE + !elementTypeOid :: Oid <- Oid . fromIntegral <$> Parser.takeInt32BE let !elementColInfo = FieldInfo elementTypeOid Nothing encodingContext when (ndim > 1) $ fail $ "TODO: No support for multi-dimensional arrays in Hpgsql. Got array with ndim=" ++ show ndim if ndim == 0 then pure mempty else do - !dim_i :: Int <- fromIntegral <$> int32Parser - !_lb_i <- int32Parser + !dim_i :: Int <- fromIntegral <$> Parser.takeInt32BE + !_lb_i <- Parser.takeInt32BE unless (elementParser.allowedPgTypes elementColInfo) $ fail $ "Array contains elements of type OID " ++ show elementTypeOid ++ " but decoder does not handle that type" replicateFunction dim_i $ do - size :: Int <- fromIntegral <$> int32Parser + size :: Int <- fromIntegral <$> Parser.takeInt32BE elementBs <- if size == (-1) then pure Nothing else Just <$> Parser.take size case elementParser.fieldValueDecoder elementColInfo elementBs of Left err -> fail $ "Error parsing array element: " ++ show err diff --git a/hpgsql/src/Hpgsql/Encoding/BinarySerializer.hs b/hpgsql/src/Hpgsql/Encoding/BinarySerializer.hs new file mode 100644 index 0000000..a21dd27 --- /dev/null +++ b/hpgsql/src/Hpgsql/Encoding/BinarySerializer.hs @@ -0,0 +1,165 @@ +{-# LANGUAGE BinaryLiterals #-} +{-# LANGUAGE CPP #-} + +-- | +-- A replacement for libraries like cereal or binary. +-- In our tests, this is ~6.0% faster than cereal, and it also +-- (or by virtue of) allocates ~13% less memory in some of our benchmarks. +-- And it also means one fewer dependency. +-- The caveat is that this module makes unaligned memory access. For the target +-- CPU architectures of this library, this should be fine. +module Hpgsql.Encoding.BinarySerializer + ( decodeInt16BE, + decodeInt32BE, + decodeInt64BE, + decodeWord32BE, + decodeWord64BE, + encodeInt32BE, + encodeDouble, + encodeFloat, + encodeInt64BE, + encodeInt16BE, + encodePgBoolean, + decodeDataRow, + ) +where + +import Data.ByteString (ByteString) +import qualified Data.ByteString.Internal as InternalBS +import Data.Int (Int16, Int32, Int64) +import Prelude hiding (encodeFloat) +#if WORDS_BIGENDIAN +import Data.Word (Word16, Word32, Word64) +#else +import Data.Word (Word16, Word32, Word64, byteSwap16, byteSwap32, byteSwap64) +#endif +import Data.Bits (Bits (unsafeShiftR)) +import qualified Data.ByteString as BS +import Data.Coerce (coerce) +import Data.Maybe (fromMaybe) +import Foreign (Storable (..), peek, (.&.)) +import Foreign.ForeignPtr (withForeignPtr) +import GHC.Float (castDoubleToWord64, castFloatToWord32) +import System.IO.Unsafe (unsafeDupablePerformIO) + +fromBigEndian32 :: Word32 -> Word32 +#if WORDS_BIGENDIAN +fromBigEndian32 = Prelude.id +#else +fromBigEndian32 = byteSwap32 +#endif + +fromBigEndian64 :: Word64 -> Word64 +#if WORDS_BIGENDIAN +fromBigEndian64 = Prelude.id +#else +fromBigEndian64 = byteSwap64 +#endif + +fromBigEndian16 :: Word16 -> Word16 +#if WORDS_BIGENDIAN +fromBigEndian16 = Prelude.id +#else +fromBigEndian16 = byteSwap16 +#endif + +{-# INLINE unsafeDecodeWord #-} +unsafeDecodeWord :: (Storable a) => ByteString -> Int -> (a -> a) -> Either String a +unsafeDecodeWord (InternalBS.BS bytesPtr len) minLen endianConvert = + if len >= minLen + then + -- A bang (strictness) in `decodedWord` makes our benchmarks allocate more memory and run slower! + let decodedWord = endianConvert $ unsafeDupablePerformIO $ withForeignPtr bytesPtr $ \ptr -> peek (coerce ptr) + in Right decodedWord + else Left "Less than enough bytes to decode" + +{-# INLINE unsafeEncodeWord #-} +unsafeEncodeWord :: (Storable a) => a -> (a -> a) -> Int -> ByteString +unsafeEncodeWord n endianConvert len = + InternalBS.unsafeCreate len $ \bufferPtr -> + poke (coerce bufferPtr) $ endianConvert n + +{-# INLINE decodeInt16BE #-} +decodeInt16BE :: ByteString -> Either String Int16 +decodeInt16BE bs = fromIntegral <$> unsafeDecodeWord bs 2 fromBigEndian16 + +{-# INLINE encodeInt16BE #-} +encodeInt16BE :: Int16 -> ByteString +encodeInt16BE n = unsafeEncodeWord (fromIntegral n) fromBigEndian16 2 + +{-# INLINE decodeWord32BE #-} +decodeWord32BE :: ByteString -> Either String Word32 +decodeWord32BE bs = unsafeDecodeWord bs 4 fromBigEndian32 + +{-# INLINE decodeWord64BE #-} +decodeWord64BE :: ByteString -> Either String Word64 +decodeWord64BE bs = unsafeDecodeWord bs 8 fromBigEndian64 + +{-# INLINE decodeInt32BE #-} +decodeInt32BE :: ByteString -> Either String Int32 +decodeInt32BE bs = fromIntegral <$> unsafeDecodeWord bs 4 fromBigEndian32 + +{-# INLINE encodeInt32BE #-} +encodeInt32BE :: Int32 -> ByteString +encodeInt32BE n = unsafeEncodeWord (fromIntegral n) fromBigEndian32 4 + +{-# INLINE decodeInt64BE #-} +decodeInt64BE :: ByteString -> Either String Int64 +decodeInt64BE bs = fromIntegral <$> unsafeDecodeWord bs 8 fromBigEndian64 + +{-# INLINE encodeInt64BE #-} +encodeInt64BE :: Int64 -> ByteString +encodeInt64BE n = unsafeEncodeWord (fromIntegral n) fromBigEndian64 8 + +{-# INLINE encodeFloat #-} +encodeFloat :: Float -> ByteString +encodeFloat n = unsafeEncodeWord (castFloatToWord32 n) fromBigEndian32 4 + +{-# INLINE encodeDouble #-} +encodeDouble :: Double -> ByteString +encodeDouble n = unsafeEncodeWord (castDoubleToWord64 n) fromBigEndian64 8 + +{-# INLINE encodePgBoolean #-} +encodePgBoolean :: Bool -> ByteString +encodePgBoolean v = if v then "\SOH" else "\NUL" + +{-# INLINE decodeDataRow #-} + +-- | A super specialized decoder to decode a postgres DataRow message +-- more quickly than a naive implementation. +-- Returns first the parsed DataRow (only column sizes and values) and second +-- the left-unparsed original bytestring. +decodeDataRow :: ByteString -> Either String (ByteString, ByteString) +decodeDataRow bs@(InternalBS.BS _bytesPtr len) = + -- We have a fast path when rows are at least 8 bytes long (should be the case + -- for all but 0-column query results or bytestring chunks "cut in the middle of the message") + -- by playing with bitwise operations. + -- Whether this is worth keeping is sort of questionable. It's complex + -- (even if I think it's safe and well tested) and reduces runtime of one of + -- our benchmarks by 2% compared to not having it. + case unsafeDecodeWord bs 8 fromBigEndian64 of + Right (w64 :: Word64) -> + -- After fromBigEndian64, the Word64 has bytes in big-endian order: + -- byte 0 (msg type) in MSB, bytes 1-4 (length) next, bytes 5-6 (col count), byte 7 in LSB. + let msgIdentByte64 = w64 .&. 0b11111111_00000000_00000000_00000000_00000000_00000000_00000000_00000000 + lenFullMsg = flip unsafeShiftR 24 $ w64 .&. 0b00000000_11111111_11111111_11111111_11111111_00000000_00000000_00000000 + letterD :: Word64 = 0b01000100_00000000_00000000_00000000_00000000_00000000_00000000_00000000 + in if msgIdentByte64 == letterD + then + toResult (fromIntegral lenFullMsg) + else Left "Not a DataRow (Word64 bits decoding path)" + Left _ -> + -- It is possible the DataRow has length less than 8 bytes, so + -- we still have to try to parse that. + if len >= 5 + then + let (InternalBS.w2c -> msgIdentChar, lenbs) = fromMaybe (error "impossible") $ BS.uncons bs + lenFullMsg = fromIntegral $ either error id (decodeInt32BE lenbs) + in if msgIdentChar == 'D' + then toResult lenFullMsg + else Left "Not a DataRow" + else Left "Less than enough bytes to decode a DataRow" + where + toResult lenFullMsg + | len >= 1 + lenFullMsg = let (a, rest) = BS.splitAt (1 + lenFullMsg) bs in Right (BS.drop 7 a, rest) + | otherwise = Left "Less than enough bytes to decode a full DataRow" diff --git a/hpgsql/src/Hpgsql/Internal.hs b/hpgsql/src/Hpgsql/Internal.hs index 8512c49..71d6de0 100644 --- a/hpgsql/src/Hpgsql/Internal.hs +++ b/hpgsql/src/Hpgsql/Internal.hs @@ -115,7 +115,6 @@ import qualified Control.Concurrent.STM as STM import Control.Exception.Safe (Exception (..), MonadThrow, SomeException, bracket, bracketOnError, finally, handleJust, mask, mask_, onException, throw, toException, tryJust) import Control.Monad (forM, forM_, join, unless, void, when) import Data.ByteString (ByteString) -import qualified Data.ByteString as BS import Data.ByteString.Internal (w2c) import qualified Data.ByteString.Lazy as LBS import Data.Data (Proxy (..)) @@ -127,7 +126,6 @@ import Data.List.NonEmpty (NonEmpty (..)) import qualified Data.List.NonEmpty as NE import qualified Data.Map.Strict as Map import Data.Maybe (fromMaybe, isNothing, mapMaybe) -import qualified Data.Serialize as Cereal import qualified Data.Set as Set import Data.Text (Text) import qualified Data.Text as Text @@ -138,6 +136,7 @@ import GHC.Conc (ThreadStatus (..), threadStatus) import Hpgsql.Base import qualified Hpgsql.Builder as Builder import Hpgsql.Encoding (FieldInfo (..), FromPgRow (..), RowDecoder (..), RowEncoder (..), ToPgRow (..)) +import qualified Hpgsql.Encoding.BinarySerializer as BinSer import Hpgsql.Encoding.RowDecoderMonadic (ConversionState (..), RowDecoderMonadic (..)) import Hpgsql.InternalTypes (BindComplete (..), CommandComplete (..), ConnectOpts (..), ConnectionString (..), CopyInResponse (..), CopyQueryState (..), DataRow (..), Either3 (..), EncodingContext (..), ErrorDetail (..), ErrorResponse (..), HPgConnection (..), InternalConnectionState (..), IrrecoverableHpgsqlError (..), NoData (..), NotificationResponse (..), ParseComplete (..), Pipeline (..), PostgresError (..), Query (..), QueryId (..), QueryProtocol (..), QueryState (..), ReadyForQuery (..), ResetConnectionOpts (..), ResponseMsg (..), ResponseMsgsReceived (..), RowDescription (..), SingleQuery (..), TransactionStatus (..), WeakThreadId (..), mkMutex, queryToByteString, throwIrrecoverableError) import Hpgsql.Locking (getMyWeakThreadId, withMutex) @@ -228,7 +227,8 @@ defaultConnectOpts = ConnectOpts { killedThreadPollIntervalMs = 500, cancellationRequestResendIntervalMs = 500, - fillTypeInfoCache = True + fillTypeInfoCache = True, + recvChunkSize = 16000 } data InternalConnectOrCancelRequest a where @@ -534,7 +534,7 @@ receiveNextMsgGeneric conn@HPgConnection {socket, recvBuffer} receiveWhat = do (initialBuf, initialBufLen) <- receiveUntilBufferHasAtLeast 5 let charAndLength = LBS.take 5 initialBuf let (w2c -> msgIdentChar, lenbs) = fromMaybe (error "impossible") $ LBS.uncons charAndLength - lenLeftToFetch :: Int64 = fromIntegral $ either error id (Cereal.decodeLazy @Int32 lenbs) - 4 + lenLeftToFetch :: Int64 = fromIntegral $ either error id (BinSer.decodeInt32BE $ LBS.toStrict lenbs) - 4 fullMessageLen = 5 + lenLeftToFetch (nowBuf, _nowBufLen) <- if initialBufLen >= fullMessageLen then pure (initialBuf, initialBufLen) else receiveUntilBufferHasAtLeast fullMessageLen let restOfMsg = LBS.drop 5 $ LBS.take fullMessageLen nowBuf @@ -595,18 +595,7 @@ receiveNextMsgGeneric conn@HPgConnection {socket, recvBuffer} receiveWhat = do fmap (bufferWithoutMsg,) $ Just <$> STM.atomically (f (Right msg)) Nothing -> handleUnexpectedMsg (f . Left) - -- Sadly we have to repeat the parsing of a DataRow message here, when it already - -- exists in the FromPgMessage instance and in the body of this function. Maybe - -- we can improve this later. - customDataRowParser = do - charAndLength <- Parser.take 5 - let (w2c -> msgIdentChar, lenbs) = fromMaybe (error "impossible") $ BS.uncons charAndLength - lenLeftToFetch :: Int = fromIntegral $ either error id (Cereal.decode @Int32 lenbs) - 4 - if msgIdentChar == 'D' - then do - rowColumnData <- BS.drop 2 <$> Parser.take lenLeftToFetch - pure $ DataRow rowColumnData - else fail "Not a DataRow" + customDataRowParser = DataRow <$> Parser.takeDataRow -- \| Appends into the internal buffer by reading from the socket -- until the buffer has at least N bytes. @@ -622,7 +611,7 @@ receiveNextMsgGeneric conn@HPgConnection {socket, recvBuffer} receiveWhat = do -- or an exception is thrown when receiving. mask $ \restore -> rethrowAsIrrecoverable $ do restore $ socketWaitRead socket - someBytes <- timeDebugNonBlockingOperation "recv" $ recvNonBlocking socket (max 16000 $ fromIntegral $ minBytesNecessary - nBytesInBuffer) + someBytes <- timeDebugNonBlockingOperation "recv" $ recvNonBlocking socket (max conn.connOpts.recvChunkSize $ fromIntegral $ minBytesNecessary - nBytesInBuffer) atomicWriteIORef recvBuffer (currentBuffer <> LBS.fromStrict someBytes) receiveUntilBufferHasAtLeast minBytesNecessary diff --git a/hpgsql/src/Hpgsql/InternalTypes.hs b/hpgsql/src/Hpgsql/InternalTypes.hs index 7234954..7b9534e 100644 --- a/hpgsql/src/Hpgsql/InternalTypes.hs +++ b/hpgsql/src/Hpgsql/InternalTypes.hs @@ -259,7 +259,7 @@ data ConnectOpts = ConnectOpts -- and you want resume using the connection and cannot wait ~500ms until Hpgsql realizes -- it's fine to do so. -- You probably don't need to worry about this or tune it. - killedThreadPollIntervalMs :: Int, + killedThreadPollIntervalMs :: !Int, -- | How long in ms Hpgsql will wait before re-sending a cancellation request -- while draining orphaned queries (queries from dead threads). The default is 500ms, -- and this is only relevant if you plan on interrupting your queries with @@ -268,14 +268,19 @@ data ConnectOpts = ConnectOpts -- It is not recommend setting this below 100ms, because orphaned query draining -- alternates with resending cancellation requests, so if this is too low it is possible -- that draining never finishes, leading to a form of livelock. - cancellationRequestResendIntervalMs :: Int, + cancellationRequestResendIntervalMs :: !Int, -- | Immediately after connecting, run a query to fetch all types -- from the `pg_type` table. This makes them available in FromPgField -- instances. -- The default is True. You should only set it to False if you really -- know what you're doing, because class instances of custom types -- can stop working. - fillTypeInfoCache :: Bool + fillTypeInfoCache :: !Bool, + -- | The minimum amount of bytes to ask for when receiving from the socket. + -- Note that Hpgsql's internal buffer may grow beyond this to accommodate + -- larger result rows. + -- The default is 16000. + recvChunkSize :: !Int } data ErrorDetail diff --git a/hpgsql/src/Hpgsql/Msgs.hs b/hpgsql/src/Hpgsql/Msgs.hs index d7fc76a..9e1f6cf 100644 --- a/hpgsql/src/Hpgsql/Msgs.hs +++ b/hpgsql/src/Hpgsql/Msgs.hs @@ -18,12 +18,12 @@ import Data.Int (Int16, Int32) import Data.Map.Strict (Map) import qualified Data.Map.Strict as Map import Data.Maybe (fromMaybe, mapMaybe) -import qualified Data.Serialize as Cereal import Data.Text (Text) import Data.Text.Encoding (decodeASCII, decodeUtf8, encodeUtf8) import Data.Word (Word8) import Hpgsql.Builder (BinaryField, Builder, builderLength) import qualified Hpgsql.Builder as Builder +import qualified Hpgsql.Encoding.BinarySerializer as BinSer import Hpgsql.InternalTypes (BindComplete (..), CommandComplete (..), CopyInResponse (..), DataRow (..), ErrorDetail (..), ErrorResponse (..), NoData (..), NotificationResponse (..), ParseComplete (..), ReadyForQuery (..), RowDescription (..), TransactionStatus (..)) import Hpgsql.ScramSHA256 (ScramClientFinalMessage (..), ScramServerFirstMessage (..)) import Hpgsql.TypeInfo (Oid (..)) @@ -54,7 +54,7 @@ colParser = do colName <- nulTerminatedCStringParser -- Column name as C string void $ Parsec.take (4 + 2) -- TODO: OIDs are unsigned integers! Try `select (-1)::oid` to see. Change to UInt32 somehow - typOid <- either fail pure . Cereal.decode @Int32 =<< Parsec.take 4 + typOid <- either fail pure . BinSer.decodeInt32BE =<< Parsec.take 4 void $ Parsec.take (2 + 4 + 2) pure (colName, Oid (fromIntegral typOid)) @@ -138,7 +138,7 @@ data Terminate = Terminate instance FromPgMessage AuthenticationResponse where msgParser = PgMsgParser $ \c restOfMsg -> case c of - 'R' -> case first (Cereal.decodeLazy @Int32) $ LBS.splitAt 4 restOfMsg of + 'R' -> case first (BinSer.decodeInt32BE . LBS.toStrict) $ LBS.splitAt 4 restOfMsg of (Right 0, _) -> Just $ AuthenticationResponse AuthOk (Right 2, _) -> Just $ AuthenticationResponse AuthKerberosV5 (Right 3, _) -> Just $ AuthenticationResponse AuthCleartextPassword @@ -155,7 +155,7 @@ instance FromPgMessage AuthenticationResponse where instance FromPgMessage BackendKeyData where msgParser = PgMsgParser $ \c (LBS.splitAt 4 -> (pidBS, backendSecretKey)) -> case c of - 'K' -> case Cereal.decodeLazy @Int32 pidBS of + 'K' -> case BinSer.decodeInt32BE $ LBS.toStrict pidBS of Right pid -> Just $ BackendKeyData {backendPid = pid, backendSecretKey = LBS.toStrict backendSecretKey} Left _ -> Nothing _ -> Nothing @@ -359,7 +359,7 @@ instance FromPgMessage RowDescription where if c == 'T' then let (numColsBS, colContents) = LBS.splitAt 2 restOfMsg - numCols = either error id $ Cereal.decodeLazy @Int16 numColsBS + numCols = either error id $ BinSer.decodeInt16BE $ LBS.toStrict numColsBS allColOidsParser :: Parsec.Parser [(Text, Oid)] allColOidsParser = replicateM (fromIntegral numCols) colParser in case LazyParsec.parseOnly (allColOidsParser <* Parsec.endOfInput) colContents of @@ -406,7 +406,7 @@ instance FromPgMessage NotificationResponse where then Nothing else let (notifierPidBs, channelNameAndPayload) = LBS.splitAt 4 restOfMsg - notifierPid = either error id $ Cereal.decodeLazy @Int32 notifierPidBs + notifierPid = either error id $ BinSer.decodeInt32BE $ LBS.toStrict notifierPidBs in case LazyParsec.parseOnly ((NotificationResponse notifierPid <$> nulTerminatedCStringParser <*> nulTerminatedCStringParser) <* Parsec.endOfInput) channelNameAndPayload of diff --git a/hpgsql/src/Hpgsql/Networking.hs b/hpgsql/src/Hpgsql/Networking.hs index b7f265d..29232c2 100644 --- a/hpgsql/src/Hpgsql/Networking.hs +++ b/hpgsql/src/Hpgsql/Networking.hs @@ -34,11 +34,11 @@ socketWaitRead socket = withFdSocket socket (threadWaitRead . fromIntegral) socketWaitWrite :: Socket -> IO () socketWaitWrite socket = withFdSocket socket (threadWaitWrite . fromIntegral) -recvNonBlocking :: Socket -> CSize -> IO ByteString -recvNonBlocking s nbytes = withFdSocket s $ \fd -> createAndTrim (fromIntegral nbytes) $ \buffer -> do +recvNonBlocking :: Socket -> Int -> IO ByteString +recvNonBlocking s nbytes = withFdSocket s $ \fd -> createAndTrim nbytes $ \buffer -> do -- Largely copied from https://hackage-content.haskell.org/package/network-3.2.8.0/docs/src/Network.Socket.Buffer.html#recvBufNoWait and other functions from the network library, -- but then modified to our needs. - r <- c_recv fd (castPtr buffer) nbytes 0 {-flags-} + r <- c_recv fd (castPtr buffer) (fromIntegral nbytes) 0 {-flags-} if r >= 0 then do -- putStrLn $ "Asked for " ++ show nbytes ++ ", got " ++ show r diff --git a/hpgsql/src/Hpgsql/SimpleParser.hs b/hpgsql/src/Hpgsql/SimpleParser.hs index 929e753..b8f1c0a 100644 --- a/hpgsql/src/Hpgsql/SimpleParser.hs +++ b/hpgsql/src/Hpgsql/SimpleParser.hs @@ -14,11 +14,17 @@ module Hpgsql.SimpleParser match, parseMany, matchLeftUnconsumed, + takeInt16BE, + takeInt32BE, + takeInt64BE, + takeDataRow, ) where import Data.ByteString (ByteString) import qualified Data.ByteString as BS +import Data.Int (Int16, Int32, Int64) +import qualified Hpgsql.Encoding.BinarySerializer as BinSer import Prelude hiding (take) data ParseResult a @@ -87,6 +93,36 @@ take n = Parser $ \bs kf ks -> ks mempty bs {-# INLINE take #-} +{-# INLINE takeInt16BE #-} +takeInt16BE :: Parser Int16 +takeInt16BE = Parser $ \bs kf ks -> + case BinSer.decodeInt16BE bs of + Left err -> kf err + Right v -> ks v (BS.drop 2 bs) + +{-# INLINE takeInt32BE #-} +takeInt32BE :: Parser Int32 +takeInt32BE = Parser $ \bs kf ks -> + case BinSer.decodeInt32BE bs of + Left err -> kf err + Right v -> ks v (BS.drop 4 bs) + +{-# INLINE takeInt64BE #-} +takeInt64BE :: Parser Int64 +takeInt64BE = Parser $ \bs kf ks -> + case BinSer.decodeInt64BE bs of + Left err -> kf err + Right v -> ks v (BS.drop 8 bs) + +{-# INLINE takeDataRow #-} + +-- | A specialized parser to parse a postgres DataRow. +takeDataRow :: Parser ByteString +takeDataRow = Parser $ \bs kf ks -> + case BinSer.decodeDataRow bs of + Left err -> kf err + Right (thisDataRow, rest) -> ks thisDataRow rest + parseMany :: Parser a -> Parser [a] parseMany p = Parser $ \bs' _kf ks -> let (vs, rest) = go bs' in ks vs rest where