50 lines
830 B
Go
50 lines
830 B
Go
package device
|
|
|
|
import (
|
|
"net"
|
|
)
|
|
|
|
func (d *Device) AllocateStreamID() uint32 {
|
|
d.streamsMu.Lock()
|
|
defer d.streamsMu.Unlock()
|
|
|
|
id := d.nextID
|
|
d.nextID++
|
|
return id
|
|
}
|
|
|
|
func (d *Device) AddStream(id uint32, conn net.Conn) chan struct{} {
|
|
d.streamsMu.Lock()
|
|
defer d.streamsMu.Unlock()
|
|
|
|
ch := make(chan struct{})
|
|
d.streams[id] = conn
|
|
d.streamDone[id] = ch
|
|
return ch
|
|
}
|
|
|
|
func (d *Device) RemoveStream(id uint32) {
|
|
d.streamsMu.Lock()
|
|
defer d.streamsMu.Unlock()
|
|
|
|
if c, ok := d.streams[id]; ok {
|
|
_ = c.Close()
|
|
delete(d.streams, id)
|
|
}
|
|
if ch, ok := d.streamDone[id]; ok {
|
|
close(ch)
|
|
delete(d.streamDone, id)
|
|
}
|
|
}
|
|
|
|
func (d *Device) GetClient(id uint32) (net.Conn, bool) {
|
|
d.streamsMu.Lock()
|
|
defer d.streamsMu.Unlock()
|
|
c, ok := d.streams[id]
|
|
return c, ok
|
|
}
|
|
|
|
func (d *Device) Done() <-chan struct{} {
|
|
return d.done
|
|
}
|