Add ability to cleaned close voice websocket and udp connections

This commit is contained in:
Bruce Marriner 2016-02-18 21:41:37 -06:00
parent b38745945b
commit 548d61c406

182
voice.go
View file

@ -46,6 +46,9 @@ type Voice struct {
guildID string
channelID string
userID string
// Used to send a close signal to goroutines
close chan struct{}
}
// ------------------------------------------------------------------------------------------------
@ -78,7 +81,8 @@ type voiceHandshakeOp struct {
// are captured.
func (v *Voice) Open() (err error) {
// TODO: How do we handle changing channels?
v.Lock()
defer v.Unlock()
// Don't open a websocket if one is already open
if v.wsConn != nil {
@ -104,45 +108,34 @@ func (v *Voice) Open() (err error) {
// Start a listening for voice websocket events
// TODO add a check here to make sure Listen worked by monitoring
// a chan or bool?
go v.wsListen()
v.close = make(chan struct{})
go v.wsListen(v.wsConn, v.close)
return
}
// Close closes the voice connection
func (v *Voice) Close() {
if v.UDPConn != nil {
err := v.UDPConn.Close()
if err != nil {
fmt.Println("error closing udp connection: ", err)
}
}
if v.wsConn != nil {
err := v.wsConn.Close()
if err != nil {
fmt.Println("error closing websocket connection: ", err)
}
}
}
// wsListen listens on the voice websocket for messages and passes them
// to the voice event handler. This is automatically called by the Open func
func (v *Voice) wsListen() {
func (v *Voice) wsListen(wsConn *websocket.Conn, close <-chan struct{}) {
for {
messageType, message, err := v.wsConn.ReadMessage()
if err != nil {
// TODO: add reconnect, matching wsapi.go:listen()
// TODO: Handle this problem better.
// TODO: needs proper logging
fmt.Println("Voice Listen Error:", err)
break
return
}
// Pass received message to voice event handler
select {
case <-close:
return
default:
go v.wsEvent(messageType, message)
}
}
return
}
@ -174,7 +167,7 @@ func (v *Voice) wsEvent(messageType int, message []byte) {
}
// Start the voice websocket heartbeat to keep the connection alive
go v.wsHeartbeat(v.OP2.HeartbeatInterval)
go v.wsHeartbeat(v.wsConn, v.close, v.OP2.HeartbeatInterval)
// TODO monitor a chan/bool to verify this was successful
// Start the UDP connection
@ -186,12 +179,16 @@ func (v *Voice) wsEvent(messageType int, message []byte) {
// Start the opusSender.
// TODO: Should we allow 48000/960 values to be user defined?
if v.OpusSend == nil {
v.OpusSend = make(chan []byte, 2)
go v.opusSender(v.OpusSend, 48000, 960)
}
go v.opusSender(v.UDPConn, v.close, v.OpusSend, 48000, 960)
// Start the opusReceiver
if v.OpusRecv == nil {
v.OpusRecv = make(chan *Packet, 2)
go v.opusReceiver(v.OpusRecv)
}
go v.opusReceiver(v.UDPConn, v.close, v.OpusRecv)
return
case 3: // HEARTBEAT response
@ -224,20 +221,33 @@ type voiceHeartbeatOp struct {
Data int `json:"d"`
}
// NOTE :: When a guild voice server changes how do we shut this down
// properly, so a new connection can be setup without fuss?
//
// wsHeartbeat sends regular heartbeats to voice Discord so it knows the client
// is still connected. If you do not send these heartbeats Discord will
// disconnect the websocket connection after a few seconds.
func (v *Voice) wsHeartbeat(i time.Duration) {
func (v *Voice) wsHeartbeat(wsConn *websocket.Conn, close <-chan struct{}, i time.Duration) {
if close == nil || wsConn == nil {
return
}
var err error
ticker := time.NewTicker(i * time.Millisecond)
for {
err := v.wsConn.WriteJSON(voiceHeartbeatOp{3, int(time.Now().Unix())})
err = wsConn.WriteJSON(voiceHeartbeatOp{3, int(time.Now().Unix())})
if err != nil {
v.Ready = false
fmt.Println("wsHeartbeat send error: ", err)
return // TODO better logging
return
}
select {
case <-ticker.C:
// continue loop and send heartbeat
case <-close:
return
}
<-ticker.C
}
}
@ -297,6 +307,21 @@ type voiceUDPOp struct {
// from voice.wsEvent OP2
func (v *Voice) udpOpen() (err error) {
v.Lock()
defer v.Unlock()
if v.wsConn == nil {
return fmt.Errorf("nil voice websocket")
}
if v.close == nil {
return fmt.Errorf("nil close channel")
}
if v.endpoint == "" {
return fmt.Errorf("empty endpoint")
}
host := fmt.Sprintf("%s:%d", strings.TrimSuffix(v.endpoint, ":80"), v.OP2.Port)
addr, err := net.ResolveUDPAddr("udp", host)
if err != nil {
@ -362,7 +387,7 @@ func (v *Voice) udpOpen() (err error) {
}
// start udpKeepAlive
go v.udpKeepAlive(5 * time.Second)
go v.udpKeepAlive(v.UDPConn, v.close, 5*time.Second)
// TODO: find a way to check that it fired off okay
return
@ -370,7 +395,11 @@ func (v *Voice) udpOpen() (err error) {
// udpKeepAlive sends a udp packet to keep the udp connection open
// This is still a bit of a "proof of concept"
func (v *Voice) udpKeepAlive(i time.Duration) {
func (v *Voice) udpKeepAlive(UDPConn *net.UDPConn, close <-chan struct{}, i time.Duration) {
if UDPConn == nil || close == nil {
return
}
var err error
var sequence uint64
@ -384,34 +413,37 @@ func (v *Voice) udpKeepAlive(i time.Duration) {
binary.LittleEndian.PutUint64(packet, sequence)
sequence++
_, err = v.UDPConn.Write(packet)
_, err = UDPConn.Write(packet)
if err != nil {
fmt.Println("udpKeepAlive udp write error : ", err)
return
}
<-ticker.C
select {
case <-ticker.C:
// continue loop and send keepalive
case <-close:
return
}
}
}
// opusSender will listen on the given channel and send any
// pre-encoded opus audio to Discord. Supposedly.
func (v *Voice) opusSender(opus <-chan []byte, rate, size int) {
func (v *Voice) opusSender(UDPConn *net.UDPConn, close <-chan struct{}, opus <-chan []byte, rate, size int) {
// TODO: Better checking to prevent this from running more than
// one instance at a time.
v.Lock()
if opus == nil {
v.Unlock()
if UDPConn == nil || close == nil {
return
}
v.Unlock()
runtime.LockOSThread()
// Voice is now ready to receive audio packets
// TODO: this needs reviewed as I think there must be a better way.
v.Ready = true
defer func() { v.Ready = false }()
defer func() {
v.Ready = false
}()
var sequence uint16
var timestamp uint32
@ -426,26 +458,32 @@ func (v *Voice) opusSender(opus <-chan []byte, rate, size int) {
ticker := time.NewTicker(time.Millisecond * time.Duration(size/(rate/1000)))
for {
// Add sequence and timestamp to udpPacket
binary.BigEndian.PutUint16(udpHeader[2:], sequence)
binary.BigEndian.PutUint32(udpHeader[4:], timestamp)
// Get data from chan. If chan is closed, return.
recvbuf, ok := <-opus
if !ok {
return
}
// Add sequence and timestamp to udpPacket
binary.BigEndian.PutUint16(udpHeader[2:], sequence)
binary.BigEndian.PutUint32(udpHeader[4:], timestamp)
// Combine the UDP Header and the opus data
sendbuf := append(udpHeader, recvbuf...)
// block here until we're exactly at the right time :)
// Then send rtp audio packet to Discord over UDP
<-ticker.C
_, err := v.UDPConn.Write(sendbuf)
select {
case <-ticker.C:
// continue
case <-close:
return
}
_, err := UDPConn.Write(sendbuf)
if err != nil {
fmt.Println("error writing to udp connection: ", err)
return
}
if (sequence) == 0xFFFF {
@ -475,27 +513,29 @@ type Packet struct {
// opusReceiver listens on the UDP socket for incoming packets
// and sends them across the given channel
// NOTE :: This function may change names later.
func (v *Voice) opusReceiver(c chan *Packet) {
func (v *Voice) opusReceiver(UDPConn *net.UDPConn, close <-chan struct{}, c chan *Packet) {
// TODO: Better checking to prevent this from running more than
// one instance at a time.
v.Lock()
if c == nil {
v.Unlock()
if UDPConn == nil || close == nil {
return
}
v.Unlock()
p := Packet{}
recvbuf := make([]byte, 1024)
for {
rlen, err := v.UDPConn.Read(recvbuf)
rlen, err := UDPConn.Read(recvbuf)
if err != nil {
fmt.Println("opusReceiver UDP Read error:", err)
return
}
select {
case <-close:
return
default:
// continue loop
}
// For now, skip anything except audio.
if rlen < 12 || recvbuf[0] != 0x80 {
continue
@ -512,3 +552,33 @@ func (v *Voice) opusReceiver(c chan *Packet) {
}
}
}
// Close closes the voice ws and udp connections
func (v *Voice) Close() {
v.Lock()
defer v.Unlock()
v.Ready = false
if v.close != nil {
close(v.close)
v.close = nil
}
if v.UDPConn != nil {
err := v.UDPConn.Close()
if err != nil {
fmt.Println("error closing udp connection: ", err)
}
v.UDPConn = nil
}
if v.wsConn != nil {
err := v.wsConn.Close()
if err != nil {
fmt.Println("error closing websocket connection: ", err)
}
v.wsConn = nil
}
}