PeerManager.hs

-----------------------------------------------------------------------------
-- |
-- Module      :  Conjure.Logic.PeerManager
-- Copyright   :  (c) Lemmih 2005, 2006
-- License     :  BSD-like
--
-- Maintainer  :  lemmih@gmail.com
-- Stability   :  experimental
-- Portability :  non-portable (requires STM)
--
{-

The described algorihm is NOT yet implemented.

-}
{-

Ideology overview:

  This bittorrent implementation builds on a few assumptions:
    * High download rates are paramount.
    * We can only download from as many leechers as we upload to.
    * Seeding is a luxury. We will only do it when we can afford it.

  There are three types of uploads:
    * Regular uploads. For peers that we're downloading from.
    * Optimistic uploads. For peers that we aren't downloading from.
    * Seed uploads. For torrents that we're seeding.

  Controversial ideas:
    * Limit upload speeds. Many medium uploads are "better" than a few
      fast ones.
    * Put peers using conjure first in the seed list.

-}
{-
WARNING: THIS TEXT IS OUT DATED AND MAY BE INACCUATE.

Choke algorithm:

Factors:
  N(10) = seconds between regular (non-optimistic and non-seeding) unchokes.
  X(11) = max number of uploads.
  Z(10) = max number of uploads - K.
  G = number of intersted peers that we're downloading from.
  J = number of uploads. (including regular, optimistic and seed uploads)
  Y = number of seed uploads.
  M(30) = seconds between optimistic unchokes.
  K(1) = min number of optimistic uploads.
  L(60) = seconds to seed before selecting a new peer.
  I(30) = seconds to (optimistically) upload before choking.

Every N seconds we:
  * sort peer list by their download speed.
  * remove uninterested peers.
  * Make sure the fastes (min G Z) are unchoked (and unchoke them if they aren't).
  * Choke slower peers.
  * J is updated.
  * if J > Z -- if total uploads > max uploads
       then stop surplus seed uploads.
       else initiate (Z - J) new seed uploads.
            J is updated.
            when Z < J
                 initiate (Z - J) new optimistic uploads.

Seed uploads continues for L seconds. After that it appends its peer
to the end of the seed list and starts uploading to the first in the
seed list. If the new peer is the same as before, no choke/unchoke
action is needed.

Every M seconds we:
  * Initiate K new optimistic uploads at random.

Optimistic uploads continues for I seconds. After that they choke the
connection unless the peer has started uploaded to us.

We're intested in a peer if:
  he got a piece we want.
  the current number of downloads < the max number of downloads.

-}
-----------------------------------------------------------------------------
module Conjure.Logic.PeerManager
    ( runMainChoker )
    where

import Conjure.Types
import Conjure.Utils
import Conjure.Protocol.PWP.Types
import Conjure.STM.PeerCtrl
import Conjure.Debug()

import GHC.Conc

import Control.Monad.State
import Text.Printf()
import Data.List

import qualified Data.Set as Set
import Data.Set (Set)

data Classification
    = Idle Integer Integer     -- download age, upload age
    | Optimistic Int           -- bps
    | Upload Int               -- bps
    | DownloadUpload (Int,Int) -- (down bps, up bps)
    | Download Int Bool        -- bps, remote interest
      deriving (Eq, Show, Ord)

type ClassifiedPeer = (Classification, ConnectedPeer)


getDownloadBps :: Classification -> Int
getDownloadBps (DownloadUpload bps) = fst bps
getDownloadBps (Download bps _) = bps
getDownloadBps _ = 0

isDownload :: Classification -> Bool
isDownload (Download _ _) = True
isDownload (DownloadUpload _) = True
isDownload _ = False

isUpload :: Classification -> Bool
isUpload (Upload _) = True
isUpload (DownloadUpload _) = True
isUpload _ = False

isIdle :: Classification -> Bool
isIdle (Idle _ _) = True
isIdle _ = False

isOptimistic :: Classification -> Bool
isOptimistic (Optimistic _) = True
isOptimistic _ = False

onFst :: (a -> a -> t1) -> t -> (a, b) -> t1
onFst f _ b = f (fst b) (fst b)

orderByUploadAge :: Classification -> Classification -> Ordering
orderByUploadAge (Idle _ age1) (Idle _ age2) = compare age1 age2
orderByUploadAge    _ _                      = EQ

orderByImportance :: Classification -> Classification -> Ordering
orderByImportance = w
    where w (Idle age1 _) (Idle age2 _) = compare age2 age1 -- Old idles are worth more.
          w (Idle _ _) _ = LT
          w _ (Idle _ _) = GT
          w (Upload bps1) (Upload bps2) = compare bps1 bps2
          w (Upload _) _ = LT
          w _ (Upload _) = GT
          w (Optimistic bps1) (Optimistic bps2) = compare bps1 bps2
          w (Optimistic _) _ = LT
          w _ (Optimistic _) = GT
          w x y = compare (getDownloadBps x) (getDownloadBps y)

classifyPeer :: Integer -> ConnectedPeer -> STM (Classification, ConnectedPeer)
classifyPeer now peer
    = do c <- classifyPeer' now peer
         return (c,peer)
classifyPeer' :: Integer -> ConnectedPeer -> STM Classification
classifyPeer' now peer
    = do lChoke <- getLocalChoke peer
         rChoke <- getRemoteChoke peer
         lInterest <- getLocalInterest peer
         rInterest <- getRemoteInterest peer
         lastDownload <- getTimingAge now (cpDownloadTimings peer)
         lastUpload <- getTimingAge now (cpUploadTimings peer)
         downBps <- getTiming now 20000 (cpDownloadTimings peer)
         upBps <- getTiming now 20000 (cpUploadTimings peer)
         return $ case (lChoke, rChoke, lInterest, rInterest) of
                    (False,False,True,True)    -> DownloadUpload (downBps, upBps)
                    (_, False, True, _)        -> Download downBps rInterest
                    (False, _, _, True)
                        | lastDownload < 30000 -> Optimistic upBps
                        | otherwise            -> Upload upBps
                    _                          -> Idle lastDownload lastUpload



choke :: ConnectedPeer -> STM ()
choke connPeer
    = do sendMessage connPeer Choke
         setLocalChoke connPeer True

unchoke :: ConnectedPeer -> STM ()
unchoke connPeer
    = do sendMessage connPeer Unchoke
         setLocalChoke connPeer False


data CState =
    CState { toChoke :: Set ConnectedPeer
           , toUnchoke :: Set ConnectedPeer }

type Unchoker a = StateT CState STM a

markForChoke :: ConnectedPeer -> Unchoker ()
markForChoke peer = modify (\s -> s{toChoke = Set.insert peer (toChoke s)
                                   ,toUnchoke = Set.delete peer (toUnchoke s)})

markForUnchoke :: ConnectedPeer -> Unchoker ()
markForUnchoke peer = modify (\s -> s{toChoke = Set.delete peer (toChoke s)
                                     ,toUnchoke = Set.insert peer (toUnchoke s)})

runMainChoker :: TVar Int -- 'N'. Use an MVar?
              -> TVar [ConnectedPeer]
              -> TVar (Int, Int, Int, Int)
              -> IO ()
runMainChoker nSecs connPeersVar opts
    = forever $
      do atomically $ do st <- execStateT (mainUnchoker connPeersVar opts) (CState Set.empty Set.empty)
                         mapM_ choke (Set.toList $ toChoke st)
                         mapM_ unchoke (Set.toList $ toUnchoke st)
         secs <- atomically (readTVar nSecs)
         threadDelay $ secs * 10^(6::Int)

mainUnchoker :: TVar [ConnectedPeer]
             -> TVar (Int,Int,Int,Int)
             -> Unchoker ()
mainUnchoker peerListVar opts
    = do (maxPeers, maxDownloads, maxUploads, minOptimistic) <- lift $ readTVar opts
         peerList <- lift $ readTVar peerListVar
         now <- lift $ unsafeIOToSTM getCurrentTime
         -- Classify peers
         classified <- lift $ mapM (classifyPeer now) peerList
         maintained <- maintainDownloads classified
         trimmed    <- trimPeers maintained maxPeers maxDownloads maxUploads
         launchOptimistic trimmed now maxUploads minOptimistic
         return ()


{-
  Keep downloads running by unchoking interested peers.
-}
maintainDownloads :: [ClassifiedPeer] -> Unchoker [ClassifiedPeer]
maintainDownloads = mapM worker
    where worker (Download bps True, peer)
              = do markForUnchoke peer
                   return (DownloadUpload (bps, 0), peer)
          worker x = return x

{-
  Trim the total number of peers, the number of downloads and the number of uploads.
-}
trimPeers :: [ClassifiedPeer] -> Int -> Int -> Int -> Unchoker [ClassifiedPeer]
trimPeers classified maxPeers maxDownloads maxUploads
    = do let byImportance = reverse (sortBy (onFst orderByImportance) classified)
             (okPeers, excessPeers) = splitAt maxPeers byImportance
         -- disconnect from the excess peers
         lift $ mapM_ (disconnectPeer.snd) excessPeers
         -- Take the downloads and sort them by speed.
         let byDownload = reverse $ sort $ filter (isDownload.fst) okPeers
             (_, excessDownloads) = splitAt maxDownloads byDownload
         -- stop downloads from excess downloads
         lift $ mapM_ (disconnectPeer.snd) excessDownloads
         let okPeers' = filter (`notElem` excessDownloads) okPeers
             byUpload = reverse $ sort $ filter (isUpload.fst) okPeers'
             (_, excessUploads) = splitAt maxUploads byUpload
         -- stop uploads to excess uploads
         mapM_ (markForChoke.snd) excessUploads
         return $ filter (`notElem` excessUploads) okPeers'

{-
  Launch new optimistic uploads
-}
launchOptimistic :: [(Classification, ConnectedPeer)] -> Integer -> Int -> Int -> Unchoker ()
launchOptimistic classified now maxUploads minOptimistic
    = do -- Start new optimistic uploads if we can.
         -- We start more than 'minOptimistic' uploads if we're low on
         -- normal uploads. This is hugely beneficial when starting a new download.
         let uploads = length (filter (isUpload.fst) classified)
         idles <- lift $ fmap (sortBy (onFst orderByUploadAge)) $
                  filterM (getRemoteInterest.snd) (filter (isIdle.fst) classified)
         let initUploads = abs (maxUploads - uploads)
             currentOptimistic = length (filter (isOptimistic.fst) classified)
             newOptimistic = abs (minOptimistic - currentOptimistic)
         forM_ (take (initUploads+newOptimistic) idles) $ \(_, peer) ->
             do lift $ recordTiming (cpDownloadTimings peer) now 0 -- Mark now as start of optimistic upload.
                markForUnchoke peer
         return ()