浏览代码

reduce memory allocation in mux

Darien Raymond 8 年之前
父节点
当前提交
a1cf299848
共有 2 个文件被更改,包括 13 次插入5 次删除
  1. 3 3
      app/proxyman/mux/mux.go
  2. 10 2
      app/proxyman/mux/session.go

+ 3 - 3
app/proxyman/mux/mux.go

@@ -370,8 +370,7 @@ func (w *ServerWorker) handleStatusEnd(meta *FrameMetadata, reader io.Reader) er
 	return nil
 }
 
-func (w *ServerWorker) handleFrame(ctx context.Context, reader io.Reader) error {
-	metaReader := NewMetadataReader(reader)
+func (w *ServerWorker) handleFrame(ctx context.Context, reader io.Reader, metaReader *MetadataReader) error {
 	meta, err := metaReader.Read()
 	if err != nil {
 		return newError("failed to read metadata").Base(err)
@@ -399,6 +398,7 @@ func (w *ServerWorker) handleFrame(ctx context.Context, reader io.Reader) error
 func (w *ServerWorker) run(ctx context.Context) {
 	input := w.outboundRay.OutboundInput()
 	reader := buf.ToBytesReader(input)
+	metaReader := NewMetadataReader(reader)
 
 	defer w.sessionManager.Close()
 
@@ -407,7 +407,7 @@ func (w *ServerWorker) run(ctx context.Context) {
 		case <-ctx.Done():
 			return
 		default:
-			err := w.handleFrame(ctx, reader)
+			err := w.handleFrame(ctx, reader, metaReader)
 			if err != nil {
 				if errors.Cause(err) != io.EOF {
 					log.Trace(newError("unexpected EOF").Base(err))

+ 10 - 2
app/proxyman/mux/session.go

@@ -19,7 +19,7 @@ type SessionManager struct {
 func NewSessionManager() *SessionManager {
 	return &SessionManager{
 		count:    0,
-		sessions: make(map[uint16]*Session, 32),
+		sessions: make(map[uint16]*Session, 16),
 	}
 }
 
@@ -58,6 +58,10 @@ func (m *SessionManager) Add(s *Session) {
 	m.Lock()
 	defer m.Unlock()
 
+	if m.closed {
+		return
+	}
+
 	m.sessions[s.ID] = s
 }
 
@@ -65,6 +69,10 @@ func (m *SessionManager) Remove(id uint16) {
 	m.Lock()
 	defer m.Unlock()
 
+	if m.closed {
+		return
+	}
+
 	delete(m.sessions, id)
 }
 
@@ -111,7 +119,7 @@ func (m *SessionManager) Close() {
 		s.output.Close()
 	}
 
-	m.sessions = make(map[uint16]*Session)
+	m.sessions = nil
 }
 
 type Session struct {