From f8b915c4e2c458db092e1e1f20ee162dae818acd Mon Sep 17 00:00:00 2001
From: Alan Braithwaite <alan@cloudflare.com>
Date: Thu, 25 Aug 2016 21:50:02 +0000
Subject: [PATCH 1/2] mem: skip defer for performance

---
 mem.go | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/mem.go b/mem.go
index a024c08c..be19ef5d 100644
--- a/mem.go
+++ b/mem.go
@@ -167,22 +167,27 @@ func (m *Message) NumSegments() int64 {
 
 // Segment returns the segment with the given ID.
 func (m *Message) Segment(id SegmentID) (*Segment, error) {
-	m.mu.Lock()
-	defer m.mu.Unlock()
+	var seg *Segment
 	if isInt32Bit() && id > maxInt32 {
 		return nil, errSegment32Bit
 	}
-	if seg := m.segment(id); seg != nil {
+	m.mu.Lock()
+	if seg = m.segment(id); seg != nil {
+		m.mu.Unlock()
 		return seg, nil
 	}
 	if int64(id) >= m.Arena.NumSegments() {
+		m.mu.Unlock()
 		return nil, errSegmentOutOfBounds
 	}
 	data, err := m.Arena.Data(id)
 	if err != nil {
+		m.mu.Unlock()
 		return nil, err
 	}
-	return m.setSegment(id, data), nil
+	seg = m.setSegment(id, data)
+	m.mu.Unlock()
+	return seg, nil
 }
 
 // segment returns the segment with the given ID.
@@ -230,19 +235,23 @@ func (m *Message) setSegment(id SegmentID, data []byte) *Segment {
 // cap(seg.Data) - len(seg.Data) >= sz.
 func (m *Message) allocSegment(sz Size) (*Segment, error) {
 	m.mu.Lock()
-	defer m.mu.Unlock()
+	var seg *Segment
 	if m.segs == nil && m.firstSeg.msg != nil {
 		m.segs = make(map[SegmentID]*Segment)
 		m.segs[0] = &m.firstSeg
 	}
 	id, data, err := m.Arena.Allocate(sz, m.segs)
 	if err != nil {
+		m.mu.Unlock()
 		return nil, err
 	}
 	if isInt32Bit() && id > maxInt32 {
+		m.mu.Unlock()
 		return nil, errSegment32Bit
 	}
-	return m.setSegment(id, data), nil
+	seg = m.setSegment(id, data)
+	m.mu.Unlock()
+	return seg, nil
 }
 
 // alloc allocates sz zero-filled bytes.  It prefers using s, but may

From 97a953fd7e044f47f4b878d24e23ccfffed28b35 Mon Sep 17 00:00:00 2001
From: Ross Light <light@google.com>
Date: Sun, 4 Sep 2016 08:58:01 -0700
Subject: [PATCH 2/2] capnp: move segment code out of critical sections

---
 mem.go | 15 ++++++---------
 1 file changed, 6 insertions(+), 9 deletions(-)

diff --git a/mem.go b/mem.go
index be19ef5d..d80c4f8f 100644
--- a/mem.go
+++ b/mem.go
@@ -167,25 +167,23 @@ func (m *Message) NumSegments() int64 {
 
 // Segment returns the segment with the given ID.
 func (m *Message) Segment(id SegmentID) (*Segment, error) {
-	var seg *Segment
 	if isInt32Bit() && id > maxInt32 {
 		return nil, errSegment32Bit
 	}
+	if int64(id) >= m.Arena.NumSegments() {
+		return nil, errSegmentOutOfBounds
+	}
 	m.mu.Lock()
-	if seg = m.segment(id); seg != nil {
+	if seg := m.segment(id); seg != nil {
 		m.mu.Unlock()
 		return seg, nil
 	}
-	if int64(id) >= m.Arena.NumSegments() {
-		m.mu.Unlock()
-		return nil, errSegmentOutOfBounds
-	}
 	data, err := m.Arena.Data(id)
 	if err != nil {
 		m.mu.Unlock()
 		return nil, err
 	}
-	seg = m.setSegment(id, data)
+	seg := m.setSegment(id, data)
 	m.mu.Unlock()
 	return seg, nil
 }
@@ -235,7 +233,6 @@ func (m *Message) setSegment(id SegmentID, data []byte) *Segment {
 // cap(seg.Data) - len(seg.Data) >= sz.
 func (m *Message) allocSegment(sz Size) (*Segment, error) {
 	m.mu.Lock()
-	var seg *Segment
 	if m.segs == nil && m.firstSeg.msg != nil {
 		m.segs = make(map[SegmentID]*Segment)
 		m.segs[0] = &m.firstSeg
@@ -249,7 +246,7 @@ func (m *Message) allocSegment(sz Size) (*Segment, error) {
 		m.mu.Unlock()
 		return nil, errSegment32Bit
 	}
-	seg = m.setSegment(id, data)
+	seg := m.setSegment(id, data)
 	m.mu.Unlock()
 	return seg, nil
 }