ceph/src/jaegertracing/thrift/lib/hs/src/Thrift/Transport/Memory.hs

   1 --
   2 -- Licensed to the Apache Software Foundation (ASF) under one
   3 -- or more contributor license agreements. See the NOTICE file
   4 -- distributed with this work for additional information
   5 -- regarding copyright ownership. The ASF licenses this file
   6 -- to you under the Apache License, Version 2.0 (the
   7 -- "License"); you may not use this file except in compliance
   8 -- with the License. You may obtain a copy of the License at
   9 --
  10 --   http://www.apache.org/licenses/LICENSE-2.0
  11 --
  12 -- Unless required by applicable law or agreed to in writing,
  13 -- software distributed under the License is distributed on an
  14 -- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  15 -- KIND, either express or implied. See the License for the
  16 -- specific language governing permissions and limitations
  17 -- under the License.
  18 --
  19
  20 module Thrift.Transport.Memory
  21        ( openMemoryBuffer
  22        , MemoryBuffer(..)
  23        ) where
  24
  25 import Data.ByteString.Lazy.Builder
  26 import Data.Functor
  27 import Data.IORef
  28 import Data.Monoid
  29 import qualified Data.ByteString.Lazy as LBS
  30
  31 import Thrift.Transport
  32
  33
  34 data MemoryBuffer = MemoryBuffer {
  35   writeBuffer :: IORef Builder,
  36   readBuffer :: IORef LBS.ByteString
  37 }
  38
  39 openMemoryBuffer :: IO MemoryBuffer
  40 openMemoryBuffer = do
  41   wbuf <- newIORef mempty
  42   rbuf <- newIORef mempty
  43   return MemoryBuffer {
  44     writeBuffer = wbuf,
  45     readBuffer = rbuf
  46   }
  47
  48 instance Transport MemoryBuffer where
  49   tIsOpen = const $ return False
  50   tClose  = const $ return ()
  51   tFlush trans = do
  52     let wBuf = writeBuffer trans
  53     wb <- readIORef wBuf
  54     modifyIORef (readBuffer trans) $ \rb -> mappend rb $ toLazyByteString wb
  55     writeIORef wBuf mempty
  56
  57   tRead _ 0 = return mempty
  58   tRead trans n = do
  59     let rbuf = readBuffer trans
  60     rb <- readIORef rbuf
  61     let len = fromIntegral $ LBS.length rb
  62     if len == 0
  63       then do
  64         tFlush trans
  65         rb2 <- readIORef (readBuffer trans)
  66         if (fromIntegral $ LBS.length rb2) == 0
  67           then return mempty
  68           else tRead trans n
  69       else do
  70         let (ret, remain) = LBS.splitAt (fromIntegral n) rb
  71         writeIORef rbuf remain
  72         return ret
  73
  74   tPeek trans = (fmap fst . LBS.uncons) <$> readIORef (readBuffer trans)
  75
  76   tWrite trans v = do
  77     modifyIORef (writeBuffer trans) (<> lazyByteString v)