diff options
Diffstat (limited to '')
-rw-r--r-- | src/jaegertracing/thrift/lib/hs/src/Thrift/Transport/Memory.hs | 77 |
1 files changed, 77 insertions, 0 deletions
diff --git a/src/jaegertracing/thrift/lib/hs/src/Thrift/Transport/Memory.hs b/src/jaegertracing/thrift/lib/hs/src/Thrift/Transport/Memory.hs new file mode 100644 index 000000000..1c93af695 --- /dev/null +++ b/src/jaegertracing/thrift/lib/hs/src/Thrift/Transport/Memory.hs @@ -0,0 +1,77 @@ +-- +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. +-- + +module Thrift.Transport.Memory + ( openMemoryBuffer + , MemoryBuffer(..) + ) where + +import Data.ByteString.Lazy.Builder +import Data.Functor +import Data.IORef +import Data.Monoid +import qualified Data.ByteString.Lazy as LBS + +import Thrift.Transport + + +data MemoryBuffer = MemoryBuffer { + writeBuffer :: IORef Builder, + readBuffer :: IORef LBS.ByteString +} + +openMemoryBuffer :: IO MemoryBuffer +openMemoryBuffer = do + wbuf <- newIORef mempty + rbuf <- newIORef mempty + return MemoryBuffer { + writeBuffer = wbuf, + readBuffer = rbuf + } + +instance Transport MemoryBuffer where + tIsOpen = const $ return False + tClose = const $ return () + tFlush trans = do + let wBuf = writeBuffer trans + wb <- readIORef wBuf + modifyIORef (readBuffer trans) $ \rb -> mappend rb $ toLazyByteString wb + writeIORef wBuf mempty + + tRead _ 0 = return mempty + tRead trans n = do + let rbuf = readBuffer trans + rb <- readIORef rbuf + let len = fromIntegral $ LBS.length rb + if len == 0 + then do + tFlush trans + rb2 <- readIORef (readBuffer trans) + if (fromIntegral $ LBS.length rb2) == 0 + then return mempty + else tRead trans n + else do + let (ret, remain) = LBS.splitAt (fromIntegral n) rb + writeIORef rbuf remain + return ret + + tPeek trans = (fmap fst . LBS.uncons) <$> readIORef (readBuffer trans) + + tWrite trans v = do + modifyIORef (writeBuffer trans) (<> lazyByteString v) |