github.com/uber/kraken@v0.1.4/lib/torrent/scheduler/announcequeue/queue.go (about)

     1  // Copyright (c) 2016-2019 Uber Technologies, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  package announcequeue
    15  
    16  import (
    17  	"container/list"
    18  
    19  	"github.com/uber/kraken/core"
    20  )
    21  
    22  // Queue manages a queue of torrents waiting to announce.
    23  type Queue interface {
    24  	Next() (core.InfoHash, bool)
    25  	Add(core.InfoHash)
    26  	Ready(core.InfoHash)
    27  	Eject(core.InfoHash)
    28  }
    29  
    30  // QueueImpl is the primary implementation of Queue. QueueImpl is not thread
    31  // safe -- synchronization must be provided by clients.
    32  type QueueImpl struct {
    33  	// Main queue of torrents ready to announce.
    34  	readyQueue *list.List
    35  
    36  	// Set of torrents with pending announce requests.
    37  	pending map[core.InfoHash]bool
    38  }
    39  
    40  // New returns a new QueueImpl.
    41  func New() *QueueImpl {
    42  	return &QueueImpl{
    43  		readyQueue: list.New(),
    44  		pending:    make(map[core.InfoHash]bool),
    45  	}
    46  }
    47  
    48  // Next returns the next torrent ready to announce. After Next is called,
    49  // the returned torrent will be marked as pending and will not be appear
    50  // again in Next until Ready is called with said torrent. Second return
    51  // value is false if no torrents are ready.
    52  func (q *QueueImpl) Next() (core.InfoHash, bool) {
    53  	next := q.readyQueue.Front()
    54  	if next == nil {
    55  		return core.InfoHash{}, false
    56  	}
    57  	q.readyQueue.Remove(next)
    58  	h := next.Value.(core.InfoHash)
    59  	q.pending[h] = true
    60  	return h, true
    61  }
    62  
    63  // Add adds a torrent to the back of the queue. Behavior is undefined if called
    64  // twice on the same torrent.
    65  func (q *QueueImpl) Add(h core.InfoHash) {
    66  	q.readyQueue.PushBack(h)
    67  }
    68  
    69  // Ready places a pending torrent back in the queue. Should be called once an
    70  // announce response is received.
    71  func (q *QueueImpl) Ready(h core.InfoHash) {
    72  	if !q.pending[h] {
    73  		return
    74  	}
    75  	delete(q.pending, h)
    76  	q.readyQueue.PushBack(h)
    77  }
    78  
    79  // Eject immediately ejects h from the announce queue, preventing it from
    80  // announcing further.
    81  func (q *QueueImpl) Eject(h core.InfoHash) {
    82  	delete(q.pending, h)
    83  	for e := q.readyQueue.Front(); e != nil; e = e.Next() {
    84  		if e.Value.(core.InfoHash) == h {
    85  			q.readyQueue.Remove(e)
    86  		}
    87  	}
    88  }
    89  
    90  // DisabledQueue is a Queue which ignores all input and constantly returns that
    91  // there are no torrents in the queue. Suitable for origin peers which want to
    92  // disable announcing.
    93  type DisabledQueue struct{}
    94  
    95  // Disabled returns a new DisabledQueue.
    96  func Disabled() DisabledQueue {
    97  	return DisabledQueue{}
    98  }
    99  
   100  // Next never returns a torrent.
   101  func (q DisabledQueue) Next() (core.InfoHash, bool) { return core.InfoHash{}, false }
   102  
   103  // Add noops.
   104  func (q DisabledQueue) Add(core.InfoHash) {}
   105  
   106  // Ready noops.
   107  func (q DisabledQueue) Ready(core.InfoHash) {}
   108  
   109  // Eject noops.
   110  func (q DisabledQueue) Eject(core.InfoHash) {}