UDT: Fix the close problem: Have a regular Close() function that may be called after reading/writing is done, and a Terminate() function that can be called anytime.

udt/multiplexer.go

@@ -50,7 +50,7 @@ func (m *multiplexer) goRead() {
 		case buf = <-m.incomingData:
 		case <-m.terminationSignal:
 			if m.socket != nil {
-				m.socket.Close() // Pass the external termination signal down to the socket. This makes sure any pending reader on the socket (especially if blocking) returns with EOF.
+				m.socket.Terminate() // Pass the external termination signal down to the socket. This makes sure any pending reader on the socket (especially if blocking) returns with EOF.
 			}
 			return
 		}

udt/udt.go

@@ -23,14 +23,15 @@ type Closer interface {
 
 // The termination reason is passed on to the close function
 const (
-	TerminateReasonListenerClosed        = 1000 // Listener: The listener.Close function was called.
-	TerminateReasonLingerTimerExpired    = 1001 // Socket: The linger timer expired. Use CloseLinger to know the actual closing reason.
-	TerminateReasonConnectTimeout        = 1002 // Socket: The connection timed out when sending the initial handshake.
-	TerminateReasonRemoteSentShutdown    = 1003 // Remote peer sent a shutdown message.
-	TerminateReasonCannotProcessOutgoing = 1004 // Send: Cannot process outgoing messages. Likely closed by caller via udtConn.
-	TerminateReasonInvalidPacketIDAck    = 1005 // Send: Invalid packet ID received in ACK message.
-	TerminateReasonInvalidPacketIDNak    = 1006 // Send: Invalid packet ID received in NAK message.
-	TerminateReasonCorruptPacketNak      = 1007 // Send: Invalid NAK packet received.
+	TerminateReasonListenerClosed     = 1000 // Listener: The listener.Close function was called.
+	TerminateReasonLingerTimerExpired = 1001 // Socket: The linger timer expired. Use CloseLinger to know the actual closing reason.
+	TerminateReasonConnectTimeout     = 1002 // Socket: The connection timed out when sending the initial handshake.
+	TerminateReasonRemoteSentShutdown = 1003 // Remote peer sent a shutdown message.
+	TerminateReasonSocketClosed       = 1004 // Send: Socket closed. Called udtSocket.Close().
+	TerminateReasonInvalidPacketIDAck = 1005 // Send: Invalid packet ID received in ACK message.
+	TerminateReasonInvalidPacketIDNak = 1006 // Send: Invalid packet ID received in NAK message.
+	TerminateReasonCorruptPacketNak   = 1007 // Send: Invalid NAK packet received.
+	TerminateReasonSignal             = 1008 // Send: Terminate signal. Called udtSocket.Terminate().
 )
 
 // DialUDT establishes an outbound UDT connection using the existing provided packet connection. It creates a UDT client.

udt/udtsocket.go

@@ -291,8 +291,8 @@ func (s *udtSocket) Write(p []byte) (n int, err error) {
 // Close closes the connection.
 // Any blocked Read or Write operations will be unblocked.
 // Write operations will be permitted to send (initial packets)
-// Read operations will return an error
-// (required for net.Conn implementation)
+// Read operations will return an error // (required for net.Conn implementation).
+// Note: Do not simultaneously call Close() and Write(). To close while the socket is still in use, use Terminate().
 func (s *udtSocket) Close() error {
 	s.closeMutex.Lock()
 	defer s.closeMutex.Unlock()
@@ -303,6 +303,22 @@ func (s *udtSocket) Close() error {
 
 	s.isClosed = true
 
+	close(s.messageOut)
+	return nil
+}
+
+// Terminate terminates the connection immediately. Unlike Close, it does not permit any reading/writing.
+// If the connection should be ordinarily closed (after reading/writing) use Close().
+func (s *udtSocket) Terminate() error {
+	s.closeMutex.Lock()
+	defer s.closeMutex.Unlock()
+
+	if s.isClosed || !s.isOpen() {
+		return nil // already closed
+	}
+
+	s.isClosed = true
+
 	close(s.terminateSignal)
 	return nil
 }

udt/udtsocket_send.go

@@ -142,7 +142,14 @@ func (s *udtSocketSend) goSendEvent() {
 
 		// wait for a channel to fire
 		select {
-		case msg := <-messageOut: // nil if we can't process outgoing messages right now, which means it will not be selected
+		case msg, ok := <-messageOut: // nil if we can't process outgoing messages right now, which means it will not be selected
+			// new message outgoing
+			if !ok {
+				s.sendPacket <- &packet.ShutdownPacket{}
+				s.shutdownEvent <- shutdownMessage{sockState: sockStateClosed, permitLinger: !s.socket.isServer, reason: TerminateReasonSocketClosed}
+				return
+			}
+
 			msg.content = s.fillDataToMTU(msg.content, messageOut) // a trick to fill up the packet immediately with data (stream only)
 
 			s.processDataMsg(msg.content, msg.tim, msg.ttl, true)
@@ -169,7 +176,7 @@ func (s *udtSocketSend) goSendEvent() {
 
 		case <-s.socket.terminateSignal:
 			s.sendPacket <- &packet.ShutdownPacket{}
-			s.shutdownEvent <- shutdownMessage{sockState: sockStateClosed, permitLinger: !s.socket.isServer, reason: TerminateReasonCannotProcessOutgoing}
+			s.shutdownEvent <- shutdownMessage{sockState: sockStateClosed, permitLinger: false, reason: TerminateReasonSignal}
 			return
 
 		case <-s.resendDataTimer: