Optimise the multicast to use one session in the tunnel

tunnel/session.go

@@ -30,8 +30,6 @@ type session struct {
 	send chan *message
 	// recv chan
 	recv chan *message
-	// wait until we have a connection
-	wait chan bool
 	// if the discovery worked
 	discovered bool
 	// if the session was accepted
@@ -109,6 +107,29 @@ func (s *session) newMessage(typ string) *message {
 	}
 }
 
+func (s *session) sendMsg(msg *message) error {
+	select {
+	case <-s.closed:
+		return io.EOF
+	case s.send <- msg:
+		return nil
+	}
+}
+
+func (s *session) wait(msg *message) error {
+	// wait for an error response
+	select {
+	case err := <-msg.errChan:
+		if err != nil {
+			return err
+		}
+	case <-s.closed:
+		return io.EOF
+	}
+
+	return nil
+}
+
 // waitFor waits for the message type required until the timeout specified
 func (s *session) waitFor(msgType string, timeout time.Duration) (*message, error) {
 	now := time.Now()
@@ -144,20 +165,32 @@ func (s *session) waitFor(msgType string, timeout time.Duration) (*message, erro
 	}
 }
 
-// Discover attempts to discover the link for a specific channel
+// Discover attempts to discover the link for a specific channel.
+// This is only used by the tunnel.Dial when first connecting.
 func (s *session) Discover() error {
 	// create a new discovery message for this channel
 	msg := s.newMessage("discover")
+	// broadcast the message to all links
 	msg.mode = Broadcast
+	// its an outbound connection since we're dialling
 	msg.outbound = true
+	// don't set the link since we don't know where it is
 	msg.link = ""
 
-	// send the discovery message
-	s.send <- msg
+	// if multicast then set that as session
+	if s.mode == Multicast {
+		msg.session = "multicast"
+	}
+
+	// send discover message
+	if err := s.sendMsg(msg); err != nil {
+		return err
+	}
 
 	// set time now
 	now := time.Now()
 
+	// after strips down the dial timeout
 	after := func() time.Duration {
 		d := time.Since(now)
 		// dial timeout minus time since
@@ -168,6 +201,7 @@ func (s *session) Discover() error {
 		return wait
 	}
 
+	// the discover message is sent out, now
 	// wait to hear back about the sent message
 	select {
 	case <-time.After(after()):
@@ -178,27 +212,16 @@ func (s *session) Discover() error {
 		}
 	}
 
-	var err error
-
-	// set a new dialTimeout
-	dialTimeout := after()
-
-	// set a shorter delay for multicast
-	if s.mode != Unicast {
-		// shorten this
-		dialTimeout = time.Millisecond * 500
-	}
-
-	// wait for announce
-	_, err = s.waitFor("announce", dialTimeout)
-
-	// if its multicast just go ahead because this is best effort
+	// bail early if its not unicast
+	// we don't need to wait for the announce
 	if s.mode != Unicast {
 		s.discovered = true
 		s.accepted = true
 		return nil
 	}
 
+	// wait for announce
+	_, err := s.waitFor("announce", after())
 	if err != nil {
 		return err
 	}
@@ -210,30 +233,22 @@ func (s *session) Discover() error {
 }
 
 // Open will fire the open message for the session. This is called by the dialler.
+// This is to indicate that we want to create a new session.
 func (s *session) Open() error {
 	// create a new message
 	msg := s.newMessage("open")
 
 	// send open message
-	s.send <- msg
+	if err := s.sendMsg(msg); err != nil {
+		return err
+	}
 
 	// wait for an error response for send
-	select {
-	case err := <-msg.errChan:
-		if err != nil {
-			return err
-		}
-	case <-s.closed:
-		return io.EOF
+	if err := s.wait(msg); err != nil {
+		return err
 	}
 
-	// don't wait on multicast/broadcast
-	if s.mode == Multicast {
-		s.accepted = true
-		return nil
-	}
-
-	// now wait for the accept
+	// now wait for the accept message to be returned
 	msg, err := s.waitFor("accept", s.timeout)
 	if err != nil {
 		return err
@@ -252,32 +267,16 @@ func (s *session) Accept() error {
 	msg := s.newMessage("accept")
 
 	// send the accept message
-	select {
-	case <-s.closed:
-		return io.EOF
-	case s.send <- msg:
-		// no op here
-	}
-
-	// don't wait on multicast/broadcast
-	if s.mode == Multicast {
-		return nil
+	if err := s.sendMsg(msg); err != nil {
+		return err
 	}
 
 	// wait for send response
-	select {
-	case err := <-s.errChan:
-		if err != nil {
-			return err
-		}
-	case <-s.closed:
-		return io.EOF
-	}
-
-	return nil
+	return s.wait(msg)
 }
 
-// Announce sends an announcement to notify that this session exists. This is primarily used by the listener.
+// Announce sends an announcement to notify that this session exists.
+// This is primarily used by the listener.
 func (s *session) Announce() error {
 	msg := s.newMessage("announce")
 	// we don't need an error back
@@ -287,23 +286,12 @@ func (s *session) Announce() error {
 	// we don't need the link
 	msg.link = ""
 
-	select {
-	case s.send <- msg:
-		return nil
-	case <-s.closed:
-		return io.EOF
-	}
+	// send announce message
+	return s.sendMsg(msg)
 }
 
 // Send is used to send a message
 func (s *session) Send(m *transport.Message) error {
-	select {
-	case <-s.closed:
-		return io.EOF
-	default:
-		// no op
-	}
-
 	// encrypt the transport message payload
 	body, err := Encrypt(m.Body, s.token+s.channel+s.session)
 	if err != nil {
@@ -335,21 +323,19 @@ func (s *session) Send(m *transport.Message) error {
 	msg.data = data
 
 	// if multicast don't set the link
-	if s.mode == Multicast {
+	if s.mode != Unicast {
 		msg.link = ""
 	}
 
 	log.Tracef("Appending %+v to send backlog", msg)
+
 	// send the actual message
-	s.send <- msg
+	if err := s.sendMsg(msg); err != nil {
+		return err
+	}
 
 	// wait for an error response
-	select {
-	case err := <-msg.errChan:
-		return err
-	case <-s.closed:
-		return io.EOF
-	}
+	return s.wait(msg)
 }
 
 // Recv is used to receive a message
@@ -413,6 +399,11 @@ func (s *session) Close() error {
 	default:
 		close(s.closed)
 
+		// don't send close on multicast
+		if s.mode != Unicast {
+			return nil
+		}
+
 		// append to backlog
 		msg := s.newMessage("close")
 		// no error response on close
@@ -421,7 +412,7 @@ func (s *session) Close() error {
 		// send the close message
 		select {
 		case s.send <- msg:
-		default:
+		case <-time.After(time.Millisecond * 10):
 		}
 	}