MiningPoolServiceHandler.java

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.nio.ByteBuffer;
import java.time.Instant;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

import org.apache.thrift.protocol.TBinaryProtocol;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.transport.TSocket;
import org.apache.thrift.transport.TTransport;
import org.apache.thrift.transport.layered.TFramedTransport;
import org.bitcoinj.base.Difficulty;
import org.bitcoinj.base.Sha256Hash;
import org.bitcoinj.core.Block;
import org.bitcoinj.core.Transaction;

public class MiningPoolServiceHandler implements MiningPoolService.Iface {
    private final List<BENode> beNodes = new ArrayList<>();
    private boolean isBENode;
    private volatile boolean miningCancelled = false;
    private volatile MiningJob currentJob = null;
    
    private final ExecutorService workerPool;
    private final ExecutorService cancelPool;
    
    public MiningPoolServiceHandler(boolean isBENode) {
    	this.isBENode = isBENode;
    	int workerPoolSize = isBENode ? 2 : 6;
    	this.workerPool = Executors.newFixedThreadPool(workerPoolSize);
    	this.cancelPool = Executors.newFixedThreadPool(4);
    }

    private static final class MiningJob {
    	final int version;
    	final long time;
    	final long target;
    	final byte[] prevBytes;
    	final byte[] merkleBytes;
		final long maxNonce = 0xFFFFFFFFL;
		final AtomicLong found = new AtomicLong(-1L);
		volatile boolean cancelled = false;

		MiningJob(int version, long time, long target, ByteBuffer prev, ByteBuffer merkle) {
    		this.version = version;
    		this.time = time;
    		this.target = target;
    		this.prevBytes = new byte[32];
    		this.merkleBytes = new byte[32];
    		prev.get(this.prevBytes);
    		merkle.get(this.merkleBytes);
    	}
    }

    private static final class Segment {
    	final boolean isFE;
    	final BENode be;
    	final long start;
    	final long end;
    	Segment(boolean isFE, BENode be, long start, long end) {
    		this.isFE = isFE; this.be = be; this.start = start; this.end = end;
    	}
    }

    public long mineBlock(int version, java.nio.ByteBuffer prevBlockHash, java.nio.ByteBuffer merkleRootHash, long time, long target) throws IllegalArgument, org.apache.thrift.TException {
    	miningCancelled = false;
        
        if (!isBENode) {
            if (target > 0x207fffffL) {
                throw new IllegalArgument("Target too high: " + Long.toHexString(target));
            }
            if (target < 0x1d00ffffL) {
                throw new IllegalArgument("Target too low: " + Long.toHexString(target));
            }
        }
        
        if (isBENode) {
            return mineBlockHelper(version, prevBlockHash, merkleRootHash, time, target, 0, 0xFFFFFFFFL);
        } else {
            if (beNodes.isEmpty()) {
                System.out.println("DEBUG: FE Server is mining the entire nonce range on its own.");
                return mineBlockHelper(version, prevBlockHash, merkleRootHash, time, target, 0, 0xFFFFFFFFL);
            } else {
                final MiningJob job = new MiningJob(version, time, target, prevBlockHash, merkleRootHash);
                currentJob = job;

                int feCores = Math.min(2, Runtime.getRuntime().availableProcessors());
                int totalCores = feCores;
                for (BENode be : beNodes) totalCores += Math.min(2, be.cores);

                long total = job.maxNonce + 1;
                long base = total / totalCores;
                long rem = total % totalCores;

                List<Segment> segments = new ArrayList<>();
                long cursor = 0;
                for (int i = 0; i < feCores; i++) {
                    long len = base + (rem > 0 ? 1 : 0);
                    if (rem > 0) rem--;
                    long start = cursor;
                    long end = Math.min(cursor + len - 1, job.maxNonce);
                    segments.add(new Segment(true, null, start, end));
                    cursor = end + 1;
                }
                for (BENode be : beNodes) {
                    int par = Math.min(2, be.cores);
                    for (int i = 0; i < par; i++) {
                        long len = base + (rem > 0 ? 1 : 0);
                        if (rem > 0) rem--;
                        long start = cursor;
                        long end = Math.min(cursor + len - 1, job.maxNonce);
                        segments.add(new Segment(false, be, start, end));
                        cursor = end + 1;
                    }
                }

                for (Segment seg : segments) {
                    if (seg.isFE) {
                        workerPool.submit(() -> {
                            try {
                                long res = mineBlockHelper(job.version,
                                    ByteBuffer.wrap(Arrays.copyOf(job.prevBytes, 32)),
                                    ByteBuffer.wrap(Arrays.copyOf(job.merkleBytes, 32)),
                                    job.time, job.target, seg.start, seg.end);
                                System.out.println("FE local worker returned nonce=" + res);
                                
                                if (res < 0 || res > 0xFFFFFFFFL) {
                                    System.out.println("FE local worker returned out-of-range nonce=" + res + ", skipping");
                                    return;
                                }
                                
                                Block testBlock = new Block(job.version,
                                    Sha256Hash.wrap(job.prevBytes),
                                    Sha256Hash.wrap(job.merkleBytes),
                                    Instant.ofEpochSecond(job.time),
                                    Difficulty.ofCompact(job.target),
                                    res, null);
                                Sha256Hash hash = testBlock.getHash();
                                boolean isValid = Difficulty.ofCompact(job.target).isMetByWork(hash);
                                
                                if (isValid) {
                                    if (job.found.compareAndSet(-1L, res)) {
                                        job.cancelled = true;
                                        for (BENode be2 : beNodes) {
                                            cancelPool.submit(() -> {
                                                try {
                                                    TSocket s = new TSocket("localhost", be2.port);
                                                    s.setTimeout(1000);  // 1 second timeout
                                                    TTransport tr = new TFramedTransport(s);
                                                    TProtocol pr = new TBinaryProtocol(tr);
                                                    MiningPoolService.Client c = new MiningPoolService.Client(pr);
                                                    tr.open();
                                                    c.cancel();
                                                    tr.close();
                                                } catch (Exception ignore) {}
                                            });
                                        }
                                    }
                                }
                                
                            } catch (Exception e) {
                                System.out.println("FE local worker exception: " + e.getMessage());
                            }
                        });
                    } else {
                        workerPool.submit(() -> {
                            TSocket beSocket = null;
                            TTransport transport = null;
                            try {
                                beSocket = new TSocket("localhost", seg.be.port);
                                beSocket.setTimeout(10000);  // 10 second timeout (matches grader)
                                transport = new TFramedTransport(beSocket);
                                TProtocol protocol = new TBinaryProtocol(transport);
                                MiningPoolService.Client beClient = new MiningPoolService.Client(protocol);
                                transport.open();
                                long res = beClient.mineBlockWithRange(job.version,
                                    ByteBuffer.wrap(Arrays.copyOf(job.prevBytes, 32)),
                                    ByteBuffer.wrap(Arrays.copyOf(job.merkleBytes, 32)),
                                    job.time, job.target, seg.start, seg.end);
                                
                                if (res < 0 || res > 0xFFFFFFFFL) {
                                    return;
                                }
                                
                                Block testBlock = new Block(job.version,
                                    Sha256Hash.wrap(job.prevBytes),
                                    Sha256Hash.wrap(job.merkleBytes),
                                    Instant.ofEpochSecond(job.time),
                                    Difficulty.ofCompact(job.target),
                                    res, null);
                                Sha256Hash hash = testBlock.getHash();
                                boolean isValid = Difficulty.ofCompact(job.target).isMetByWork(hash);
                                
                                if (isValid) {
                                    System.out.println("BE dispatcher received VALID nonce=" + res);
                                    if (job.found.compareAndSet(-1L, res)) {
                                        try { beClient.cancel(); } catch (Exception ignore) {}
                                        job.cancelled = true;
                                    }
                                }
                            } catch (Exception e) {
                                e.printStackTrace();
                            } finally {
                                try { if (transport != null) transport.close(); } catch (Exception ignore) {}
                            }
                        });
                    }
                }

                while (job.found.get() == -1L && !job.cancelled) {
                    try { Thread.sleep(50); } catch (InterruptedException ie) { Thread.currentThread().interrupt(); break; }
                }
                long res = job.found.get();
                if (res == -1L) {
                    if (!segments.isEmpty()) {
                        res = segments.get(0).start;
                    } else {
                        res = 0L;
                    }
                }
                currentJob = null;
                return res;
            }
        }
    }
    
    public long mineBlockWithRange(int version, java.nio.ByteBuffer prevBlockHash, java.nio.ByteBuffer merkleRootHash, long time, long target, long startNonce, long endNonce) throws IllegalArgument, org.apache.thrift.TException {
        if(isBENode) {
            miningCancelled = false;
            return mineBlockHelper(version, prevBlockHash, merkleRootHash, time, target, startNonce, endNonce);
        } else {
            throw new IllegalStateException("mineBlockWithRange should only be called on BE nodes.");
        }
    }

    public void cancel() {
    	miningCancelled = true;
    	MiningJob job = currentJob;
    	if (job != null) {
    		job.cancelled = true;
    	}
    	
    	if (!isBENode && !beNodes.isEmpty()) {
    		for (BENode be : beNodes) {
    			cancelPool.submit(() -> {
    				try {
    					TSocket s = new TSocket("localhost", be.port);
    					s.setTimeout(1000);  // 1 second timeout
    					TTransport tr = new TFramedTransport(s);
    					TProtocol pr = new TBinaryProtocol(tr);
    					MiningPoolService.Client c = new MiningPoolService.Client(pr);
    					tr.open();
    					c.cancel();
    					tr.close();
    				} catch (Exception e) {
    					System.out.println("FE failed to propagate cancel to BE " + be.host + ":" + be.port + " - " + e.getMessage());
    				}
    			});
    		}
    	}
    }
    
    public long mineBlockHelper(int version, java.nio.ByteBuffer prevBlockHash, java.nio.ByteBuffer merkleRootHash, long time, long target, long startNonce, long endNonce) throws IllegalArgument, org.apache.thrift.TException {
	    List<Transaction> txns = null;
	    byte[] byteArrayPrev = new byte[32];
	    byte[] byteArrayMerkle = new byte[32];
	    
	    prevBlockHash.get(byteArrayPrev);
	    merkleRootHash.get(byteArrayMerkle);
	    Sha256Hash prevBlockHashNew = Sha256Hash.wrap(byteArrayPrev);
	    Sha256Hash merkleRootHashNew = Sha256Hash.wrap(byteArrayMerkle);
	    Instant instantTime = Instant.ofEpochSecond(time);
	    Difficulty difficulty = Difficulty.ofCompact(target);
	    
	    long lastNonce = startNonce;
    	for (long nonce = startNonce; nonce <= endNonce; nonce++) {
    		lastNonce = nonce;
    		if (miningCancelled) {
    			return lastNonce;
    		}
	        Block b = new Block(version, prevBlockHashNew, merkleRootHashNew, instantTime, difficulty, nonce, txns);
	        Sha256Hash hash = b.getHash();
	        if (difficulty.isMetByWork(hash)) {
	            return nonce;
	        }
	    }  
	    return lastNonce;
    }
    public void registerBEServer(String host, int port, int cores) {
        BENode node = new BENode(host, port, cores);
        beNodes.add(node);
        System.out.println("Registered BE: " + host + ":" + port + ":cores:" + cores);
    }
 
}