3 files changed, 36 insertions, 23 deletions
diff --git a/simulation/config/config.go b/simulation/config/config.go
index 797145c..d22f4f1 100644
--- a/simulation/config/config.go
+++ b/simulation/config/config.go
@@ -56,14 +56,18 @@ type Node struct {
 	Changes   []Change
 }
 
+// LatencyModel for ths simulation.
+type LatencyModel struct {
+	Mean  float64
+	Sigma float64
+}
+
 // Networking config.
 type Networking struct {
 	Type       test.NetworkType
 	PeerServer string
-
-	Mean          float64
-	Sigma         float64
-	LossRateValue float64
+	Direct     LatencyModel
+	Gossip     LatencyModel
 }
 
 // Scheduler Settings.
@@ -129,11 +133,16 @@ func GenerateDefault(path string) error {
 			MaxBlock: math.MaxUint64,
 		},
 		Networking: Networking{
-			Type:          test.NetworkTypeTCPLocal,
-			PeerServer:    "127.0.0.1",
-			Mean:          100,
-			Sigma:         10,
-			LossRateValue: 0,
+			Type:       test.NetworkTypeTCPLocal,
+			PeerServer: "127.0.0.1",
+			Direct: LatencyModel{
+				Mean:  100,
+				Sigma: 10,
+			},
+			Gossip: LatencyModel{
+				Mean:  300,
+				Sigma: 25,
+			},
 		},
 		Scheduler: Scheduler{
 			WorkerNum: 2,
diff --git a/simulation/node.go b/simulation/node.go
index dc4a725..026db66 100644
--- a/simulation/node.go
+++ b/simulation/node.go
@@ -70,18 +70,19 @@ type node struct {
 func newNode(prvKey crypto.PrivateKey, logger common.Logger,
 	cfg config.Config) *node {
 	pubKey := prvKey.PublicKey()
-	netModule := test.NewNetwork(
-		pubKey,
-		&test.NormalLatencyModel{
-			Mean:  cfg.Networking.Mean,
-			Sigma: cfg.Networking.Sigma,
+	netModule := test.NewNetwork(pubKey, test.NetworkConfig{
+		Type:       cfg.Networking.Type,
+		PeerServer: cfg.Networking.PeerServer,
+		PeerPort:   peerPort,
+		DirectLatency: &test.NormalLatencyModel{
+			Mean:  cfg.Networking.Direct.Mean,
+			Sigma: cfg.Networking.Direct.Sigma,
 		},
-		test.NewDefaultMarshaller(&jsonMarshaller{}),
-		test.NetworkConfig{
-			Type:       cfg.Networking.Type,
-			PeerServer: cfg.Networking.PeerServer,
-			PeerPort:   peerPort,
-		})
+		GossipLatency: &test.NormalLatencyModel{
+			Mean:  cfg.Networking.Gossip.Mean,
+			Sigma: cfg.Networking.Gossip.Sigma,
+		},
+		Marshaller: test.NewDefaultMarshaller(&jsonMarshaller{})})
 	id := types.NewNodeID(pubKey)
 	dbInst, err := db.NewMemBackedDB(id.String() + ".db")
 	if err != nil {
diff --git a/simulation/peer-server.go b/simulation/peer-server.go
index 14a825a..69ed029 100644
--- a/simulation/peer-server.go
+++ b/simulation/peer-server.go
@@ -90,7 +90,8 @@ func (p *PeerServer) handleBlockList(id types.NodeID, blocks *BlockList) {
 		}
 		p.verifiedLen += uint64(length)
 		if p.verifiedLen >= p.cfg.Node.MaxBlock {
-			if err := p.trans.Broadcast(ntfShutdown); err != nil {
+			if err := p.trans.Broadcast(
+				p.peers, &test.FixedLatencyModel{}, ntfShutdown); err != nil {
 				panic(err)
 			}
 		}
@@ -199,7 +200,8 @@ func (p *PeerServer) Run() {
 	}
 	// Cache peers' info.
 	for _, pubKey := range p.trans.Peers() {
-		p.peers[types.NewNodeID(pubKey)] = struct{}{}
+		nID := types.NewNodeID(pubKey)
+		p.peers[nID] = struct{}{}
 	}
 	// Pick a mater node to execute pending config changes.
 	for nID := range p.peers {
@@ -225,7 +227,8 @@ func (p *PeerServer) Run() {
 			break
 		}
 	}
-	if err := p.trans.Broadcast(ntfReady); err != nil {
+	if err := p.trans.Broadcast(
+		p.peers, &test.FixedLatencyModel{}, ntfReady); err != nil {
 		panic(err)
 	}
 	log.Println("Simulation is ready to go with", len(p.peers), "nodes")