Move where out messages are created.

connection.go

@@ -86,9 +86,10 @@ type Connection struct {
 // State that is maintained for each in-flight op. This is stuffed into the
 // context that the user uses to reply to the op.
 type opState struct {
-	inMsg *buffer.InMessage
-	op    interface{}
-	opID  uint32 // For logging
+	inMsg  *buffer.InMessage
+	outMsg *buffer.OutMessage
+	op     interface{}
+	opID   uint32 // For logging
 }
 
 // Create a connection wrapping the supplied file descriptor connected to the
@@ -370,14 +371,14 @@ func (c *Connection) ReadOp() (ctx context.Context, op interface{}, err error) {
 	// Keep going until we find a request we know how to convert.
 	for {
 		// Read the next message from the kernel.
-		var m *buffer.InMessage
-		m, err = c.readMessage()
+		var inMsg *buffer.InMessage
+		inMsg, err = c.readMessage()
 		if err != nil {
 			return
 		}
 
 		// Convert the message to an op.
-		op, err = convertInMessage(m, c.protocol)
+		op, err = convertInMessage(inMsg, c.protocol)
 		if err != nil {
 			err = fmt.Errorf("convertInMessage: %v", err)
 			return
@@ -395,9 +396,12 @@ func (c *Connection) ReadOp() (ctx context.Context, op interface{}, err error) {
 			continue
 		}
 
+		// Allocate an output message up front, to be used later when replying.
+		outMsg := c.getOutMessage()
+
 		// Set up a context that remembers information about this op.
-		ctx = c.beginOp(m.Header().Opcode, m.Header().Unique)
-		ctx = context.WithValue(ctx, contextKey, opState{m, op, opID})
+		ctx = c.beginOp(inMsg.Header().Opcode, inMsg.Header().Unique)
+		ctx = context.WithValue(ctx, contextKey, opState{inMsg, outMsg, op, opID})
 
 		// Special case: responding to statfs is required to make mounting work on
 		// OS X. We don't currently expose the capability for the file system to
@@ -426,14 +430,16 @@ func (c *Connection) Reply(ctx context.Context, opErr error) {
 	}
 
 	op := state.op
-	m := state.inMsg
+	inMsg := state.inMsg
+	outMsg := state.outMsg
 	opID := state.opID
 
-	// Make sure we destroy the message when we're done.
-	defer c.putInMessage(m)
+	// Make sure we destroy the messages when we're done.
+	defer c.putInMessage(inMsg)
+	defer c.putOutMessage(outMsg)
 
 	// Clean up state for this op.
-	c.finishOp(m.Header().Opcode, m.Header().Unique)
+	c.finishOp(inMsg.Header().Opcode, inMsg.Header().Unique)
 
 	// Debug logging
 	if c.debugLogger != nil {
@@ -450,8 +456,7 @@ func (c *Connection) Reply(ctx context.Context, opErr error) {
 	}
 
 	// Send the reply to the kernel, if one is required.
-	outMsg := c.getOutMessage()
-	noResponse := c.kernelResponse(outMsg, m.Header().Unique, op, opErr)
+	noResponse := c.kernelResponse(outMsg, inMsg.Header().Unique, op, opErr)
 
 	if !noResponse {
 		err := c.writeMessage(outMsg.Bytes())
@@ -459,8 +464,6 @@ func (c *Connection) Reply(ctx context.Context, opErr error) {
 			c.errorLogger.Printf("writeMessage: %v", err)
 		}
 	}
-
-	c.putOutMessage(outMsg)
 }
 
 // Close the connection. Must not be called until operations that were read