diff --git a/agents/base_agent.go b/agents/base_agent.go index 1c40df9..ef3501e 100644 --- a/agents/base_agent.go +++ b/agents/base_agent.go @@ -80,6 +80,13 @@ func (c *ConnectionAgents) Get(name string) Agent { return c.agents[name] } +func (c *ConnectionAgents) Stop(names... string) { + for _, n := range names { + c.Get(n).Stop() + c.Get(n).Join() + } +} + func (c *ConnectionAgents) StopAll() { for _, a := range c.agents { a.Stop() diff --git a/scenarii/ack_only.go b/scenarii/ack_only.go index 42ff52b..1e86432 100644 --- a/scenarii/ack_only.go +++ b/scenarii/ack_only.go @@ -25,8 +25,7 @@ func (s *AckOnlyScenario) Run(conn *qt.Connection, trace *qt.Trace, preferredUrl return } defer connAgents.CloseConnection(false, 0, "") - connAgents.Get("AckAgent").Stop() - connAgents.Get("AckAgent").Join() + connAgents.Stop("AckAgent") incPackets := conn.IncomingPackets.RegisterNewChan(1000) diff --git a/scenarii/connection_migration.go b/scenarii/connection_migration.go index 199dbf1..321868d 100644 --- a/scenarii/connection_migration.go +++ b/scenarii/connection_migration.go @@ -30,10 +30,7 @@ func (s *ConnectionMigrationScenario) Run(conn *qt.Connection, trace *qt.Trace, <-time.NewTimer(3 * time.Second).C // Wait some time before migrating - connAgents.Get("SocketAgent").Stop() - connAgents.Get("SendingAgent").Stop() - connAgents.Get("SocketAgent").Join() - connAgents.Get("SendingAgent").Join() + connAgents.Stop("SocketAgent", "SendingAgent") newUdpConn, err := qt.EstablishUDPConnection(conn.Host) if err != nil { diff --git a/scenarii/zero_rtt.go b/scenarii/zero_rtt.go index 255c88f..f0b0f49 100644 --- a/scenarii/zero_rtt.go +++ b/scenarii/zero_rtt.go @@ -66,8 +66,7 @@ func (s *ZeroRTTScenario) Run(conn *qt.Connection, trace *qt.Trace, preferredUrl } connAgents = agents.AttachAgentsToConnection(conn, agents.GetDefaultAgents()...) - connAgents.Get("RecoveryAgent").Stop() - connAgents.Get("RecoveryAgent").Join() + connAgents.Stop("RecoveryAgent") handshakeAgent := &agents.HandshakeAgent{TLSAgent: connAgents.Get("TLSAgent").(*agents.TLSAgent), SocketAgent: connAgents.Get("SocketAgent").(*agents.SocketAgent)} connAgents.Add(handshakeAgent) defer connAgents.CloseConnection(false, 0, "")