2 files changed, 57 insertions, 11 deletions

diff --git a/examples/dhtd.hs b/examples/dhtd.hs
index fed2976f..57ee8deb 100644
--- a/examples/dhtd.hs
+++ b/examples/dhtd.hs
@@ -673,20 +673,23 @@ clientSession s@Session{..} sock cnum h = do
             -> cmd0 $ do
                 sessions' <- atomically $ readTVar sessions :: IO [PerSession]
                 let sessionsReport = mapM showPerSession sessions'
-                    headers = ["Key", "NextMsg", "Handled","Unhandled"]
+                    headers = ["Key", "NextMsg", "Dropped","Handled","Unhandled"]
                     showPerSession (PerSession
                                         { perSessionMsgs = msgQ
                                         , perSessionPublicKey = pubKey
                                         , perSessionAddr = sockAddr
                                         , perSessionNumVar = msgNumVar
+                                        , perSessionDropCount = dropCntVar
                                         }) = do
                         num <- atomically (readTVar msgNumVar)
+                        dropped <- atomically (readTVar dropCntVar)
                         as <- atomically (packetQueueViewList msgQ)
                         let (h,u) = partition (fst . snd) as
                             countHandled   = length h
                             countUnhandled = length u
                         return [ show (Tox.key2id pubKey) -- "Key"
                                , show num                 -- "NextMsg"
+                               , show dropped             -- "Dropped"
                                , show countHandled        -- "Handled"
                                , show countUnhandled      -- "Unhandled"
                                ]
@@ -1248,7 +1251,8 @@ announceToxJabberPeer echan laddr saddr pingflag tsrc tsnk
 data PerSession = PerSession { perSessionMsgs :: PacketQueue (Bool{-Handled?-},Tox.CryptoMessage)
                              , perSessionPublicKey :: PublicKey
                              , perSessionAddr :: SockAddr
-                             , perSessionNumVar :: TVar Word32
+                             , perSessionNumVar    :: TVar Word32
+                             , perSessionDropCount :: TVar Word32
                              }
 
 main :: IO ()
@@ -1572,12 +1576,14 @@ main = runResourceT $ liftBaseWith $ \resT -> do
             -- allsessionsMap <- atomically $ readTVar (netCryptoSessions netCryptoSessionsState)
             let sockAddr = Tox.ncSockAddr netcrypto
                 pubKey   = Tox.ncTheirPublicKey netcrypto
-            msgQ <- atomically (Data.PacketQueue.new 10 0 :: STM (PacketQueue (Bool,Tox.CryptoMessage)))
+            msgQ <- atomically (Data.PacketQueue.newOverwrite 10 0 :: STM (PacketQueue (Bool,Tox.CryptoMessage)))
             msgNumVar <- atomically (newTVar 0)
+            dropCntVar <- atomically (newTVar 0)
             let perSession = PerSession { perSessionMsgs = msgQ
                                         , perSessionPublicKey = pubKey
                                         , perSessionAddr = sockAddr
                                         , perSessionNumVar = msgNumVar
+                                        , perSessionDropCount = dropCntVar
                                         }
             atomically $ modifyTVar' sessions (perSession:)
             tmchan <- atomically newTMChan
@@ -1599,8 +1605,11 @@ main = runResourceT $ liftBaseWith $ \resT -> do
                 handleIncoming mTyp session cm = do
                     atomically $ do
                         num <- readTVar msgNumVar
-                        enqueue msgQ num (False,cm)
+                        (wraps,offset) <- enqueue msgQ num (False,cm)
+                        capacity <- getCapacity msgQ
+                        let dropped = wraps * capacity + offset
                         modifyTVar' msgNumVar (+1)
+                        writeTVar dropCntVar dropped
                     atomically $ writeTMChan tmchan cm -- (Tox.bufferData cd)
                     return Nothing
             atomically $ writeTVar (Tox.ncUnrecognizedHook netcrypto) handleIncoming
diff --git a/src/Data/PacketQueue.hs b/src/Data/PacketQueue.hs
index a4c99cab..4f0f04e3 100644
--- a/src/Data/PacketQueue.hs
+++ b/src/Data/PacketQueue.hs
@@ -6,8 +6,11 @@
 {-# LANGUAGE FlexibleContexts #-}
 module Data.PacketQueue
     ( PacketQueue
+    , getCapacity
     , new
+    , newOverwrite
     , dequeue
+    , markButNotDequeue
     , enqueue
     , observeOutOfBand
     , PacketOutQueue
@@ -24,6 +27,7 @@ module Data.PacketQueue
 import Control.Concurrent.STM
 import Control.Concurrent.STM.TArray
 import Control.Monad
+import Control.Applicative
 import Data.Word
 import Data.Array.MArray
 import Data.Maybe
@@ -33,6 +37,7 @@ data PacketQueue a = PacketQueue
     , seqno   :: TVar Word32
     , qsize   :: Word32
     , buffend :: TVar Word32 -- on incoming, highest packet number handled + 1
+    , qOverWriteMode :: Bool
     }
 
 packetQueueViewList :: PacketQueue a -> STM [(Word32,a)]
@@ -41,6 +46,8 @@ packetQueueViewList p = do
         f (n,Just x)  = Just (n,x)
     catMaybes . map f <$> getAssocs (pktq p)
 
+getCapacity :: Applicative m => PacketQueue t -> m Word32
+getCapacity (PacketQueue { qsize }) = pure qsize
 
 -- | Create a new PacketQueue.
 new :: Word32 -- ^ Capacity of queue.
@@ -56,6 +63,24 @@ new capacity seqstart = do
         , seqno   = seqv
         , qsize   = cap
         , buffend = bufe
+        , qOverWriteMode = False
+        }
+
+-- | Create a new PacketQueue with Overwrite on Wrap.
+newOverwrite :: Word32 -- ^ Capacity of queue.
+    -> Word32 -- ^ Initial sequence number.
+    -> STM (PacketQueue a)
+newOverwrite capacity seqstart = do
+    let cap = if capacity `mod` 2 == 0 then capacity else capacity + 1
+    q <- newArray (0,cap - 1) Nothing
+    seqv <- newTVar seqstart
+    bufe <- newTVar 0
+    return PacketQueue
+        { pktq    = q
+        , seqno   = seqv
+        , qsize   = cap
+        , buffend = bufe
+        , qOverWriteMode = True
         }
 
 observeOutOfBand :: PacketQueue a -> Word32-> STM ()
@@ -77,21 +102,33 @@ dequeue PacketQueue { pktq, seqno, qsize } = do
     modifyTVar' seqno   succ
     return x
 
+-- | Like dequeue, but marks as handled rather than removing
+markButNotDequeue :: PacketQueue (Bool,a) -> STM a
+markButNotDequeue PacketQueue { pktq, seqno, qsize } = do
+    i0 <- readTVar seqno
+    let i = i0 `mod` qsize
+    (b,x) <- maybe retry return =<< readArray pktq i
+    writeArray pktq i (Just (True,x))
+    modifyTVar' seqno   succ
+    return x
+
 -- | Enqueue a packet.  Packets need not be enqueued in order as long as there
 -- is spare capacity in the queue.  If there is not, the packet will be
--- silently discarded without blocking.
+-- silently discarded without blocking. (Unless this is an Overwrite-queue,
+-- in which case, the packets will simply wrap around overwriting the old ones.)
 enqueue :: PacketQueue a -- ^ The packet queue.
         -> Word32        -- ^ Sequence number of the packet.
         -> a             -- ^ The packet.
-        -> STM ()
-enqueue PacketQueue{ pktq, seqno, qsize, buffend } no x = do
+        -> STM (Word32,Word32)
+enqueue PacketQueue{ pktq, seqno, qsize, buffend,qOverWriteMode } no x = do
     low <- readTVar seqno
     let proj = no - low
     -- Ignore packet if out of range.
-    when ( proj < qsize) $ do
-    let i = no `mod` qsize
-    writeArray pktq i (Just x)
-    modifyTVar' buffend (\be -> if be - low <= proj then no + 1 else be)
+    when ( proj < qsize || qOverWriteMode) $ do
+        let i = no `mod` qsize
+        writeArray pktq i (Just x)
+        modifyTVar' buffend (\be -> if be - low <= proj then no + 1 else be)
+    return (proj `divMod` qsize)
 
 -- lookup :: PacketQueue a -> Word32 -> STM (Maybe a)
 -- lookup PacketQueue{ pktq, seqno, qsize } no = _todo