Merge pull request #48 from abraithwaite/performance

mem.go: skip defer for performance

mem.go

@@ -168,22 +168,25 @@ 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()
 	if isInt32Bit() && id > maxInt32 {
 		return nil, errSegment32Bit
 	}
-	if seg := m.segment(id); seg != nil {
-		return seg, nil
-	}
 	if int64(id) >= m.Arena.NumSegments() {
 		return nil, errSegmentOutOfBounds
 	}
+	m.mu.Lock()
+	if seg := m.segment(id); seg != nil {
+		m.mu.Unlock()
+		return seg, nil
+	}
 	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.
@@ -231,19 +234,22 @@ 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()
 	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