Update Broadcast Stage logic

src/main/java/com/iota/iri/network/pipeline/BroadcastStage.java

@@ -3,6 +3,7 @@
 import com.iota.iri.controllers.TransactionViewModel;
 import com.iota.iri.network.NeighborRouter;
 import com.iota.iri.network.neighbor.Neighbor;
+import com.iota.iri.service.validation.TransactionSolidifier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -18,13 +19,16 @@ public class BroadcastStage implements Stage {
 
     private NeighborRouter neighborRouter;
 
+    private TransactionSolidifier transactionSolidifier;
+
     /**
      * Creates a new {@link BroadcastStage}.
      * 
      * @param neighborRouter The {@link NeighborRouter} instance to use to broadcast
      */
-    public BroadcastStage(NeighborRouter neighborRouter) {
+    public BroadcastStage(NeighborRouter neighborRouter, TransactionSolidifier transactionSolidifier) {
         this.neighborRouter = neighborRouter;
+        this.transactionSolidifier = transactionSolidifier;
     }
 
     /**
@@ -54,6 +58,15 @@ public ProcessingContext process(ProcessingContext ctx) {
             }
         }
 
+        // Check the transaction solidifier to see if there are solid transactions that need to be broadcast.
+        // If so, forward them to the BroadcastStageQueue to be processed.
+        TransactionViewModel transactionToBroadcast;
+        if((transactionToBroadcast = transactionSolidifier.getNextTxInBroadcastQueue()) != null){
+            ctx.setNextStage(TransactionProcessingPipeline.Stage.BROADCAST);
+            ctx.setPayload(new BroadcastPayload(payload.getOriginNeighbor(), transactionToBroadcast));
+            return ctx;
+        }
+
         ctx.setNextStage(TransactionProcessingPipeline.Stage.FINISH);
         return ctx;
     }

src/main/java/com/iota/iri/network/pipeline/SolidifyStage.java

@@ -71,21 +71,15 @@ public ProcessingContext process(ProcessingContext ctx){
     }
 
     private ProcessingContext broadcastTip(ProcessingContext ctx, SolidifyPayload payload) throws  Exception{
-        // First check if there is a transaction available to broadcast from the broadcast queue
-        TransactionViewModel tip = txSolidifier.getNextTxInBroadcastQueue();
+        Hash tipHash = tipsViewModel.getRandomSolidTipHash();
 
-        // If there is not a transaction available from the broadcast queue, instead try to send a solid tip
-        if (tip == null) {
-            Hash tipHash = tipsViewModel.getRandomSolidTipHash();
-
-            if (tipHash == null) {
-                ctx.setNextStage(TransactionProcessingPipeline.Stage.FINISH);
-                return ctx;
-            }
-
-            tip = fromHash(tangle, tipHash);
+        if (tipHash == null) {
+            ctx.setNextStage(TransactionProcessingPipeline.Stage.FINISH);
+            return ctx;
         }
 
+        TransactionViewModel tip = fromHash(tangle, tipHash);
+
         ctx.setNextStage(TransactionProcessingPipeline.Stage.BROADCAST);
         ctx.setPayload(new BroadcastPayload(payload.getOriginNeighbor(), tip));
 

src/main/java/com/iota/iri/network/pipeline/TransactionProcessingPipelineImpl.java

@@ -99,7 +99,7 @@ public TransactionProcessingPipelineImpl(NeighborRouter neighborRouter, NodeConf
         this.preProcessStage = new PreProcessStage(recentlySeenBytesCache);
         this.replyStage = new ReplyStage(neighborRouter, config, tangle, tipsViewModel, latestMilestoneTracker,
                 snapshotProvider, recentlySeenBytesCache);
-        this.broadcastStage = new BroadcastStage(neighborRouter);
+        this.broadcastStage = new BroadcastStage(neighborRouter, txSolidifier);
         this.validationStage = new ValidationStage(txValidator, recentlySeenBytesCache);
         this.receivedStage = new ReceivedStage(tangle, txSolidifier, snapshotProvider, transactionRequester);
         this.batchedHasher = BatchedHasherFactory.create(BatchedHasherFactory.Type.BCTCURL81, 20);

src/test/java/com/iota/iri/network/pipeline/BroadcastStageTest.java

@@ -9,6 +9,7 @@
 import java.util.HashMap;
 import java.util.Map;
 
+import com.iota.iri.service.validation.TransactionSolidifier;
 import org.junit.Rule;
 import org.junit.Test;
 import org.mockito.Mock;
@@ -21,6 +22,9 @@ public class BroadcastStageTest {
     @Rule
     public MockitoRule mockitoRule = MockitoJUnit.rule();
 
+    @Mock
+    private TransactionSolidifier transactionSolidifier;
+
     @Mock
     private NeighborRouter neighborRouter;
 
@@ -38,7 +42,7 @@ public class BroadcastStageTest {
     public void doesntGossipToOriginNeighbor() {
         Mockito.when(neighborRouter.getConnectedNeighbors()).thenReturn(neighbors);
 
-        BroadcastStage broadcastStage = new BroadcastStage(neighborRouter);
+        BroadcastStage broadcastStage = new BroadcastStage(neighborRouter, transactionSolidifier);
         TransactionViewModel tvm = new TransactionViewModel(new Transaction(), null);
         BroadcastPayload broadcastPayload = new BroadcastPayload(neighborA, tvm);
         ProcessingContext ctx = new ProcessingContext(null, broadcastPayload);
@@ -58,7 +62,7 @@ public void doesntGossipToOriginNeighbor() {
     public void gossipsToAllIfNoOriginNeighbor() {
         Mockito.when(neighborRouter.getConnectedNeighbors()).thenReturn(neighbors);
 
-        BroadcastStage broadcastStage = new BroadcastStage(neighborRouter);
+        BroadcastStage broadcastStage = new BroadcastStage(neighborRouter, transactionSolidifier);
         TransactionViewModel tvm = new TransactionViewModel(new Transaction(), null);
         BroadcastPayload broadcastPayload = new BroadcastPayload(null, tvm);
         ProcessingContext ctx = new ProcessingContext(null, broadcastPayload);