add WriteBuffers function to send a vector of slice in batch

This commit is contained in:
xtaci
2019-04-07 21:56:50 +08:00
parent 59b7c4b915
commit 7c104a5443
2 changed files with 86 additions and 0 deletions
+61
View File
@@ -304,6 +304,67 @@ func (s *UDPSession) Write(b []byte) (n int, err error) {
}
}
// WriteBuffers write a vector of byte slices to the underlying connection
func (s *UDPSession) WriteBuffers(v [][]byte) (n int, err error) {
for {
s.mu.Lock()
if s.isClosed {
s.mu.Unlock()
return 0, errors.New(errBrokenPipe)
}
if s.kcp.WaitSnd() < int(s.kcp.snd_wnd) {
for _, b := range v {
n += len(b)
for {
if len(b) <= int(s.kcp.mss) {
s.kcp.Send(b)
break
} else {
s.kcp.Send(b[:s.kcp.mss])
b = b[s.kcp.mss:]
}
}
}
if s.kcp.WaitSnd() >= int(s.kcp.snd_wnd) || !s.writeDelay {
s.kcp.flush(false)
}
s.mu.Unlock()
atomic.AddUint64(&DefaultSnmp.BytesSent, uint64(n))
return n, nil
}
var timeout *time.Timer
var c <-chan time.Time
if !s.wd.IsZero() {
if time.Now().After(s.wd) {
s.mu.Unlock()
return 0, errTimeout{}
}
delay := s.wd.Sub(time.Now())
timeout = time.NewTimer(delay)
c = timeout.C
}
s.mu.Unlock()
select {
case <-s.chWriteEvent:
case <-c:
case <-s.die:
case err = <-s.chWriteError:
if timeout != nil {
timeout.Stop()
}
return n, err
}
if timeout != nil {
timeout.Stop()
}
}
}
// Close closes the connection.
func (s *UDPSession) Close() error {
// remove current session from updater & listener(if necessary)
+25
View File
@@ -272,6 +272,31 @@ func TestSendRecv(t *testing.T) {
cli.Close()
}
func TestSendVector(t *testing.T) {
cli, err := dialEcho()
if err != nil {
panic(err)
}
cli.SetWriteDelay(false)
const N = 100
buf := make([]byte, 20)
v := make([][]byte, 2)
for i := 0; i < N; i++ {
v[0] = []byte(fmt.Sprintf("hello%v", i))
v[1] = []byte(fmt.Sprintf("world%v", i))
msg := fmt.Sprintf("hello%vworld%v", i, i)
cli.WriteBuffers(v)
if n, err := cli.Read(buf); err == nil {
if string(buf[:n]) != msg {
t.Error(string(buf[:n]), msg)
}
} else {
panic(err)
}
}
cli.Close()
}
func TestTinyBufferReceiver(t *testing.T) {
cli, err := dialTinyBufferEcho()
if err != nil {