性能进行优化

session.go

@@ -377,10 +377,31 @@ func (s *Session) keepalive() {
 	}
 }
 
+func (s *Session) bodyClone(buf []byte, p *sendReady) {
+	if buf == nil {
+		return // A nil body is ignored.
+	}
+
+	// In the event of session shutdown or connection write timeout,
+	// we need to prevent `send` from reading the body buffer after
+	// returning from this function since the caller may re-use the
+	// underlying array.
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	if p.Body == nil {
+		return // Body was already copied in `send`.
+	}
+	newBody := make([]byte, len(buf))
+	copy(newBody, buf)
+	p.Body = newBody
+}
+
 // waitForSendErr waits to send a header, checking for a potential shutdown
-func (s *Session) waitForSend(hdr header, body []byte) error {
+func (s *Session) waitForSend(hdr header, body []byte) (err error) {
 	errCh := make(chan error, 1)
-	return s.waitForSendErr(hdr, body, errCh)
+	err = s.waitForSendErr(hdr, body, errCh)
+	return
 }
 
 // waitForSendErr waits to send a header with optional data, checking for a
@@ -398,7 +419,6 @@ func (s *Session) waitForSendErr(hdr header, body []byte, errCh chan error) erro
 		}
 		timerPool.Put(t)
 	}()
-
 	ready := &sendReady{Hdr: hdr, Body: body, Err: errCh}
 	select {
 	case s.sendCh <- ready:
@@ -407,35 +427,14 @@ func (s *Session) waitForSendErr(hdr header, body []byte, errCh chan error) erro
 	case <-timer.C:
 		return ErrConnectionWriteTimeout
 	}
-
-	bodyCopy := func() {
-		if body == nil {
-			return // A nil body is ignored.
-		}
-
-		// In the event of session shutdown or connection write timeout,
-		// we need to prevent `send` from reading the body buffer after
-		// returning from this function since the caller may re-use the
-		// underlying array.
-		ready.mu.Lock()
-		defer ready.mu.Unlock()
-
-		if ready.Body == nil {
-			return // Body was already copied in `send`.
-		}
-		newBody := make([]byte, len(body))
-		copy(newBody, body)
-		ready.Body = newBody
-	}
-
 	select {
 	case err := <-errCh:
 		return err
 	case <-s.shutdownCh:
-		bodyCopy()
+		s.bodyClone(body, ready)
 		return ErrSessionShutdown
 	case <-timer.C:
-		bodyCopy()
+		s.bodyClone(body, ready)
 		return ErrConnectionWriteTimeout
 	}
 }

session_test.go

@@ -7,8 +7,11 @@ import (
 	"io"
 	"io/ioutil"
 	"net"
+	"os"
+	"path"
 	"reflect"
 	"runtime"
+	"runtime/pprof"
 	"strings"
 	"sync"
 	"testing"
@@ -388,6 +391,16 @@ func TestNonNilInterface(t *testing.T) {
 }
 
 func TestSendData_Small(t *testing.T) {
+	var (
+		count = 100000
+	)
+	filename := path.Join(os.TempDir(), "yamux.pprof")
+	t.Log(filename)
+	fp, _ := os.OpenFile(filename, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
+	defer func() {
+		fp.Close()
+	}()
+	pprof.StartCPUProfile(fp)
 	client, server := testClientServer()
 	defer client.Close()
 	defer server.Close()
@@ -407,7 +420,7 @@ func TestSendData_Small(t *testing.T) {
 		}
 
 		buf := make([]byte, 4)
-		for i := 0; i < 1000; i++ {
+		for i := 0; i < count; i++ {
 			n, err := stream.Read(buf)
 			if err != nil {
 				t.Fatalf("err: %v", err)
@@ -436,7 +449,7 @@ func TestSendData_Small(t *testing.T) {
 			t.Fatalf("bad")
 		}
 
-		for i := 0; i < 1000; i++ {
+		for i := 0; i < count; i++ {
 			n, err := stream.Write([]byte("test"))
 			if err != nil {
 				t.Fatalf("err: %v", err)
@@ -464,6 +477,7 @@ func TestSendData_Small(t *testing.T) {
 		if server.NumStreams() != 0 {
 			t.Fatalf("bad")
 		}
+		pprof.StopCPUProfile()
 		return
 	case <-time.After(time.Second):
 		panic("timeout")