diff --git a/sess.go b/sess.go index 2fdac0d..c6d5435 100644 --- a/sess.go +++ b/sess.go @@ -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) diff --git a/sess_test.go b/sess_test.go index 4fce29a..1eaa471 100644 --- a/sess_test.go +++ b/sess_test.go @@ -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 {