package bt.dht;

import bt.metainfo.Torrent;
import bt.net.Peer;
import bt.peer.PeerSource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Collection;
import java.util.HashSet;
import java.util.Queue;
import java.util.Set;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.locks.ReentrantLock;
import java.util.stream.Stream;

/**
 * @since 1.1
 */
public class DHTPeerSource implements PeerSource {

    private static final Logger LOGGER = LoggerFactory.getLogger(DHTPeerSource.class);

    private static final int MAX_PEERS_PER_BATCH = 10;

    private final Torrent torrent;
    private final DHTService dhtService;
    private final ExecutorService executor;

    private final ReentrantLock lock;
    private final AtomicReference<Future<?>> futureOptional;
    private final Queue<Peer> peers;

    DHTPeerSource(Torrent torrent, DHTService dhtService, ExecutorService executor) {
        this.torrent = torrent;
        this.dhtService = dhtService;
        this.executor = executor;

        this.lock = new ReentrantLock();
        this.futureOptional = new AtomicReference<>();
        this.peers = new LinkedBlockingQueue<>();
    }

    @Override
    public boolean update() {
        if (peers.isEmpty()) {
            schedulePeerCollection();
        }
        return !peers.isEmpty();
    }

    private void schedulePeerCollection() {
        if (lock.tryLock()) {
            try {
                if (futureOptional.get() != null) {
                    Future<?> future = futureOptional.get();
                    if (future.isDone()) {
                        try {
                            future.get();
                        } catch (InterruptedException | ExecutionException e) {
                            LOGGER.warn("Peer collection finished with exception for torrent ID: " + torrent.getTorrentId(), e);
                        }
                        futureOptional.set(null);
                    }
                }

                if (futureOptional.get() == null) {
                    futureOptional.set(executor.submit(this::collectPeers));
                    LOGGER.info("Scheduled peer collection for torrent ID: " + torrent.getTorrentId());
                }
            } finally {
                lock.unlock();
            }
        }
    }

    private void collectPeers() {
        Stream<Peer> peerStream = dhtService.getPeers(torrent);
        peerStream.forEach(peer -> {
            peers.add(peer);
            if (LOGGER.isDebugEnabled()) {
                LOGGER.debug(String.format("Collected new peer (torrent ID: %s, peer: %s)", torrent.getTorrentId(), peer));
            }
        });
        LOGGER.info("Peer collection finished for torrent ID: " + torrent.getTorrentId());
    }

    @Override
    public Collection<Peer> getPeers() {
        Set<Peer> batch = new HashSet<>((int)(MAX_PEERS_PER_BATCH / 0.75d + 1));
        int i = 0;
        Peer peer;
        while (i++ < MAX_PEERS_PER_BATCH && (peer = peers.poll()) != null) {
            batch.add(peer);
        }
        if (LOGGER.isDebugEnabled()) {
            LOGGER.debug("Collected a batch of " + batch.size() + " peers: " + batch);
        }
        return batch;
    }
}

@@ -73,16 +73,25 @@
 bt.peer.
+Scheduled
 PeerSour
@@ -173,369 +173,70 @@
 til.
-Collection;%0Aimport java.util.HashSet;%0Aimport java.util.Queue;%0Aimport java.util.Set;%0Aimport java.util.concurrent.ExecutionException;%0Aimport java.util.concurrent.ExecutorService;%0Aimport java.util.concurrent.Future;%0Aimport java.util.concurrent.LinkedBlockingQueue;%0Aimport java.util.concurrent.atomic.AtomicReference;%0Aimport java.util.concurrent.locks.ReentrantLock
+concurrent.ExecutorService;%0Aimport java.util.function.Consumer
 ;%0Aim
@@ -319,19 +319,25 @@
 rce 
-implement
+extend
 s 
+Scheduled
 Peer
@@ -481,17 +481,22 @@
 PER_
-BATCH = 1
+COLLECTION = 5
 0;%0A%0A
@@ -574,189 +574,8 @@
 ice;
-%0A    private final ExecutorService executor;%0A%0A    private final ReentrantLock lock;%0A    private final AtomicReference%3CFuture%3C?%3E%3E futureOptional;%0A    private final Queue%3CPeer%3E peers;
 %0A%0A  
@@ -654,24 +654,49 @@
 executor) %7B%0A
+        super(executor);%0A
         this
@@ -761,1368 +761,84 @@
     
-    this.executor = executor;%0A%0A        this.lock = new ReentrantLock();%0A        this.futureOptional = new AtomicReference%3C%3E();%0A        this.peers = new LinkedBlockingQueue%3C%3E();%0A    %7D%0A%0A    @Override%0A    public boolean update() %7B%0A        if (peers.isEmpty()) %7B%0A            schedulePeerCollection();%0A        %7D%0A        return !peers.isEmpty();%0A    %7D%0A%0A    private void schedulePeerCollection() %7B%0A        if (lock.tryLock()) %7B%0A            try %7B%0A                if (futureOptional.get() != null) %7B%0A                    Future%3C?%3E future = futureOptional.get();%0A                    if (future.isDone()) %7B%0A                        try %7B%0A                            future.get();%0A                        %7D catch (InterruptedException %7C ExecutionException e) %7B%0A                            LOGGER.warn(%22Peer collection finished with exception for torrent ID: %22 + torrent.getTorrentId(), e);%0A                        %7D%0A                        futureOptional.set(null);%0A                    %7D%0A                %7D%0A%0A                if (futureOptional.get() == null) %7B%0A                    futureOptional.set(executor.submit(this::collectPeers));%0A                    LOGGER.info(%22Scheduled peer collection for torrent ID: %22 + torrent.getTorrentId());%0A                %7D%0A            %7D finally %7B%0A                lock.unlock();%0A            %7D%0A        %7D%0A    %7D%0A%0A    private void collectPeers(
+%7D%0A%0A    @Override%0A    protected void collectPeers(Consumer%3CPeer%3E peerConsumer
 ) %7B%0A
@@ -899,16 +899,48 @@
 torrent)
+.limit(MAX_PEERS_PER_COLLECTION)
 ;%0A      
@@ -990,13 +990,23 @@
 peer
-s.add
+Consumer.accept
 (pee
@@ -1302,462 +1302,7 @@
    %7D
-%0A%0A    @Override%0A    public Collection%3CPeer%3E getPeers() %7B%0A        Set%3CPeer%3E batch = new HashSet%3C%3E((int)(MAX_PEERS_PER_BATCH / 0.75d + 1));%0A        int i = 0;%0A        Peer peer;%0A        while (i++ %3C MAX_PEERS_PER_BATCH && (peer = peers.poll()) != null) %7B%0A            batch.add(peer);%0A        %7D%0A        if (LOGGER.isDebugEnabled()) %7B%0A            LOGGER.debug(%22Collected a batch of %22 + batch.size() + %22 peers: %22 + batch);%0A        %7D%0A        return batch;%0A    %7D
 %0A%7D%0A

Use ScheduledPeerSource in DHT module