1 files changed, 37 insertions, 17 deletions

diff --git a/src/cpu/testers/traffic_gen/traffic_gen.hh b/src/cpu/testers/traffic_gen/traffic_gen.hh
index c013109b7..0adcf781e 100644
--- a/src/cpu/testers/traffic_gen/traffic_gen.hh
+++ b/src/cpu/testers/traffic_gen/traffic_gen.hh
@@ -42,6 +42,7 @@
 #define __CPU_TRAFFIC_GEN_TRAFFIC_GEN_HH__
 
 #include "base/hashmap.hh"
+#include "base/statistics.hh"
 #include "cpu/testers/traffic_gen/generators.hh"
 #include "mem/mem_object.hh"
 #include "mem/qport.hh"
@@ -74,18 +75,6 @@ class TrafficGen : public MemObject
     void enterState(uint32_t newState);
 
     /**
-     * Get the tick of the next event, either a new packet or a
-     * transition.
-     *
-     * @return tick of the next update event
-     */
-    Tick nextEventTick()
-    {
-        return std::min(states[currState]->nextPacketTick(),
-                        nextTransitionTick);
-    }
-
-    /**
      * Parse the config file and build the state map and
      * transition matrix.
      */
@@ -98,6 +87,12 @@ class TrafficGen : public MemObject
      */
     void update();
 
+    /**
+     * Receive a retry from the neighbouring port and attempt to
+     * resend the waiting packet.
+     */
+    void recvRetry();
+
     /** Struct to represent a probabilistic transition during parsing. */
     struct Transition {
         uint32_t from;
@@ -124,6 +119,9 @@ class TrafficGen : public MemObject
     /** Time of next transition */
     Tick nextTransitionTick;
 
+    /** Time of the next packet. */
+    Tick nextPacketTick;
+
     /** State transition matrix */
     std::vector<std::vector<double> > transitionMatrix;
 
@@ -133,31 +131,50 @@ class TrafficGen : public MemObject
     /** Map of generator states */
     m5::hash_map<uint32_t, BaseGen*> states;
 
-    /** Queued master port */
-    class TrafficGenPort : public QueuedMasterPort
+    /** Master port specialisation for the traffic generator */
+    class TrafficGenPort : public MasterPort
     {
       public:
 
-        TrafficGenPort(const std::string& name, TrafficGen& _owner)
-            : QueuedMasterPort(name, &_owner, queue), queue(_owner, *this)
+        TrafficGenPort(const std::string& name, TrafficGen& traffic_gen)
+            : MasterPort(name, &traffic_gen), trafficGen(traffic_gen)
         { }
 
       protected:
 
+        void recvRetry() { trafficGen.recvRetry(); }
+
         bool recvTimingResp(PacketPtr pkt);
 
       private:
 
-        MasterPacketQueue queue;
+        TrafficGen& trafficGen;
 
     };
 
     /** The instance of master port used by the traffic generator. */
     TrafficGenPort port;
 
+    /** Packet waiting to be sent. */
+    PacketPtr retryPkt;
+
+    /** Tick when the stalled packet was meant to be sent. */
+    Tick retryPktTick;
+
     /** Event for scheduling updates */
     EventWrapper<TrafficGen, &TrafficGen::update> updateEvent;
 
+    /** Manager to signal when drained */
+    DrainManager* drainManager;
+
+    /** Count the number of generated packets. */
+    Stats::Scalar numPackets;
+
+    /** Count the number of retries. */
+    Stats::Scalar numRetries;
+
+    /** Count the time incurred from back-pressure. */
+    Stats::Scalar retryTicks;
 
   public:
 
@@ -178,6 +195,9 @@ class TrafficGen : public MemObject
 
     void unserialize(Checkpoint* cp, const std::string& section);
 
+    /** Register statistics */
+    void regStats();
+
 };
 
 #endif //__CPU_TRAFFIC_GEN_TRAFFIC_GEN_HH__