the allocator can now be enabled per request

Other minor changes as per review comments

.travis.yml

@@ -35,12 +35,12 @@ script:
  - go test -integration -v ./...
  - go test -testserver -v ./...
  - go test -integration -testserver -v ./...
- - go test -integration -optimized-allocator -v ./...
- - go test -testserver -optimized-allocator -v ./...
- - go test -integration -testserver -optimized-allocator -v ./...
+ - go test -integration -allocator -v ./...
+ - go test -testserver -allocator -v ./...
+ - go test -integration -testserver -allocator -v ./...
  - go test -race -integration -v ./...
  - go test -race -testserver -v ./...
  - go test -race -integration -testserver -v ./...
- - go test -race -integration -optimized-allocator -v ./...
- - go test -race -testserver -optimized-allocator -v ./...
- - go test -race -integration -optimized-allocator -testserver -v ./...
+ - go test -race -integration -allocator -v ./...
+ - go test -race -testserver -allocator -v ./...
+ - go test -race -integration -allocator -testserver -v ./...

Makefile

@@ -2,16 +2,16 @@ integration:
  go test -integration -v
  go test -testserver -v
  go test -integration -testserver -v
- go test -integration -optimized-allocator -v ./...
- go test -testserver -optimized-allocator -v ./...
- go test -integration -testserver -optimized-allocator -v ./...
+ go test -integration -allocator -v ./...
+ go test -testserver -allocator -v ./...
+ go test -integration -testserver -allocator -v ./...
 
 integration_w_race:
  go test -race -integration -v
  go test -race -testserver -v
  go test -race -integration -testserver -v
- go test -race -integration -optimized-allocator -v ./...
- go test -race -testserver -optimized-allocator -v ./...
- go test -race -integration -optimized-allocator -testserver -v ./...
+ go test -race -integration -allocator -v ./...
+ go test -race -testserver -allocator -v ./...
+ go test -race -integration -allocator -testserver -v ./...
 
 

allocator.go

@@ -5,10 +5,10 @@ import (
 )
 
 type allocator struct {
+ sync.Mutex
  available [][]byte
  // map key is the request order
  used map[uint32][][]byte
- sync.Mutex
 }
 
 func newAllocator() *allocator {
@@ -27,7 +27,7 @@ func (a *allocator) GetPage(requestOrderID uint32) []byte {
 
  var result []byte
 
- // get an available page and remove it from the available ones
+ // get an available page and remove it from the available ones.
  if len(a.available) > 0 {
  truncLength := len(a.available) - 1
  result = a.available[truncLength]
@@ -52,15 +52,13 @@ func (a *allocator) ReleasePages(requestOrderID uint32) {
  a.Lock()
  defer a.Unlock()
 
- if used, ok := a.used[requestOrderID]; ok && len(used) > 0 {
+ if used := a.used[requestOrderID]; len(used) > 0 {
  a.available = append(a.available, used...)
- // this is probably useless
- a.used[requestOrderID] = nil
  }
  delete(a.used, requestOrderID)
 }
 
-// Free removes all the used and free pages.
+// Free removes all the used and available pages.
 // Call this method when the allocator is not needed anymore
 func (a *allocator) Free() {
  a.Lock()
@@ -92,8 +90,6 @@ func (a *allocator) isRequestOrderIDUsed(requestOrderID uint32) bool {
  a.Lock()
  defer a.Unlock()
 
- if _, ok := a.used[requestOrderID]; ok {
- return true
- }
- return false
+ _, ok := a.used[requestOrderID]
+ return ok
 }

allocator_test.go

@@ -8,7 +8,6 @@ import (
  "github.com/stretchr/testify/assert"
 )
 
-// I like the full flow of the test here, but probably I will be asked to split it in separate test cases
 func TestAllocator(t *testing.T) {
  allocator := newAllocator()
  // get a page for request order id 1

client.go

@@ -214,7 +214,7 @@ func (c *Client) nextID() uint32 {
 }
 
 func (c *Client) recvVersion() error {
- typ, data, err := c.recvPacket(nil, 0)
+ typ, data, err := c.recvPacket(0)
  if err != nil {
  return err
  }

conn.go

@@ -13,15 +13,17 @@ import (
 type conn struct {
  io.Reader
  io.WriteCloser
+ // this is the same allocator used in packet manager
+ alloc *allocator
  sync.Mutex // used to serialise writes to sendPacket
  // sendPacketTest is needed to replicate packet issues in testing
  sendPacketTest func(w io.Writer, m encoding.BinaryMarshaler) error
 }
 
-// the allocator and the orderID are used in server mode if AllocationModeOptimized is enabled.
-// For the client just pass nil and 0
-func (c *conn) recvPacket(alloc *allocator, orderID uint32) (uint8, []byte, error) {
- return recvPacket(c, alloc, orderID)
+// the orderID is used in server mode if the allocator is enabled.
+// For the client mode just pass 0
+func (c *conn) recvPacket(orderID uint32) (uint8, []byte, error) {
+ return recvPacket(c, c.alloc, orderID)
 }
 
 func (c *conn) sendPacket(m encoding.BinaryMarshaler) error {
@@ -78,7 +80,7 @@ func (c *clientConn) recv() error {
  c.conn.Close()
  }()
  for {
- typ, data, err := c.recvPacket(nil, 0)
+ typ, data, err := c.recvPacket(0)
  if err != nil {
  return err
  }

packet-manager.go

@@ -18,7 +18,7 @@ type packetManager struct {
  sender packetSender // connection object
  working *sync.WaitGroup
  packetCount uint32
- // it is not nil if AllocationModeOptimized is enabled
+ // it is not nil if the allocator is enabled
  alloc *allocator
 }
 
@@ -36,9 +36,6 @@ func newPktMgr(sender packetSender) *packetManager {
  sender: sender,
  working: &sync.WaitGroup{},
  }
- if enabledAllocationMode == AllocationModeOptimized {
- s.alloc = newAllocator()
- }
  go s.controller()
  return s
 }
@@ -50,7 +47,7 @@ func (s *packetManager) newOrderID() uint32 {
 }
 
 // returns the next orderID without incrementing it.
-// This is used before receiving a new packet in AllocationModeOptimized to associate
+// This is used before receiving a new packet, with the allocator enabled, to associate
 // the slice allocated for the received packet with the orderID that will be used to mark
 // the allocated slices for reuse once the request is served
 func (s *packetManager) getNextOrderID() uint32 {

packet.go

@@ -149,24 +149,24 @@ func recvPacket(r io.Reader, alloc *allocator, orderID uint32) (uint8, []byte, e
  if _, err := io.ReadFull(r, b[:4]); err != nil {
  return 0, nil, err
  }
- l, _ := unmarshalUint32(b)
- if l > maxMsgLength {
- debug("recv packet %d bytes too long", l)
+ length, _ := unmarshalUint32(b)
+ if length > maxMsgLength {
+ debug("recv packet %d bytes too long", length)
  return 0, nil, errLongPacket
  }
  if alloc == nil {
- b = make([]byte, l)
+ b = make([]byte, length)
  }
- if _, err := io.ReadFull(r, b[0:l]); err != nil {
- debug("recv packet %d bytes: err %v", l, err)
+ if _, err := io.ReadFull(r, b[:length]); err != nil {
+ debug("recv packet %d bytes: err %v", length, err)
  return 0, nil, err
  }
  if debugDumpRxPacketBytes {
- debug("recv packet: %s %d bytes %x", fxp(b[0]), l, b[1:l])
+ debug("recv packet: %s %d bytes %x", fxp(b[0]), length, b[1:length])
  } else if debugDumpRxPacket {
- debug("recv packet: %s %d bytes", fxp(b[0]), l)
+ debug("recv packet: %s %d bytes", fxp(b[0]), length)
  }
- return b[0], b[1:l], nil
+ return b[0], b[1:length], nil
 }
 
 type extensionPair struct {
@@ -593,11 +593,13 @@ func (p *sshFxpReadPacket) UnmarshalBinary(b []byte) error {
 
 func (p *sshFxpReadPacket) getDataSlice(alloc *allocator, orderID uint32) []byte {
  dataLen := clamp(p.Len, maxTxPacket)
- // we allocate a slice with a bigger capacity so we avoid a new allocation in sshFxpDataPacket.MarshalBinary
- // and in sendPacket, we need 9 bytes in MarshalBinary and 4 bytes in sendPacket.
  if alloc != nil {
+ // GetPage returns a slice with capacity = maxMsgLength this is enough to avoid new allocations in
+ // sshFxpDataPacket.MarshalBinary and sendPacket
  return alloc.GetPage(orderID)[:dataLen]
  }
+ // we allocate a slice with a bigger capacity so we avoid a new allocation in sshFxpDataPacket.MarshalBinary
+ // and in sendPacket, we need 9 bytes in MarshalBinary and 4 bytes in sendPacket.
  return make([]byte, dataLen, dataLen+9+4)
 }
 

request-server.go

@@ -33,20 +33,50 @@ type RequestServer struct {
 }
 
 // NewRequestServer creates/allocates/returns new RequestServer.
-// Normally there there will be one server per user-session.
+// Normally there will be one server per user-session.
 func NewRequestServer(rwc io.ReadWriteCloser, h Handlers) *RequestServer {
+ rs, _ := NewRequestServerWithOptions(rwc, h, nil...)
+ return rs
+}
+
+// A RequestServerOption is a function which applies configuration to a RequestServer.
+type RequestServerOption func(*RequestServer) error
+
+// WithRSAllocator enable the allocator.
+// After processing a packet we keep in memory the allocated slices
+// and we reuse them for new packets.
+// The allocator is experimental
+func WithRSAllocator() RequestServerOption {
+ return func(rs *RequestServer) error {
+ alloc := newAllocator()
+ rs.pktMgr.alloc = alloc
+ rs.conn.alloc = alloc
+ return nil
+ }
+}
+
+// NewRequestServerWithOptions creates/allocates/returns new RequestServer adding the specified options
+// If options is nil or empty this is equivalent to NewRequestServer
+func NewRequestServerWithOptions(rwc io.ReadWriteCloser, h Handlers, options ...RequestServerOption) (*RequestServer, error) {
  svrConn := &serverConn{
  conn: conn{
  Reader: rwc,
  WriteCloser: rwc,
  },
  }
- return &RequestServer{
+ rs := &RequestServer{
  serverConn: svrConn,
  Handlers: h,
  pktMgr: newPktMgr(svrConn),
  openRequests: make(map[string]*Request),
  }
+
+ for _, o := range options {
+ if err := o(rs); err != nil {
+ return nil, err
+ }
+ }
+ return rs, nil
 }
 
 // New Open packet/Request
@@ -88,6 +118,11 @@ func (rs *RequestServer) Close() error { return rs.conn.Close() }
 
 // Serve requests for user session
 func (rs *RequestServer) Serve() error {
+ defer func() {
+ if rs.pktMgr.alloc != nil {
+ rs.pktMgr.alloc.Free()
+ }
+ }()
  ctx, cancel := context.WithCancel(context.Background())
  defer cancel()
  var wg sync.WaitGroup
@@ -107,7 +142,7 @@ func (rs *RequestServer) Serve() error {
  var pktType uint8
  var pktBytes []byte
  for {
- pktType, pktBytes, err = rs.recvPacket(rs.pktMgr.alloc, rs.pktMgr.getNextOrderID())
+ pktType, pktBytes, err = rs.serverConn.recvPacket(rs.pktMgr.getNextOrderID())
  if err != nil {
  // we don't care about releasing allocated pages here, the server will quit and the allocator freed
  break
@@ -150,9 +185,6 @@ func (rs *RequestServer) Serve() error {
  delete(rs.openRequests, handle)
  req.close()
  }
- if rs.pktMgr.alloc != nil {
- rs.pktMgr.alloc.Free()
- }
 
  return err
 }

request-server_test.go

@@ -46,7 +46,12 @@ func clientRequestServerPair(t *testing.T) *csPair {
  fd, err := l.Accept()
  assert.Nil(t, err)
  handlers := InMemHandler()
- server = NewRequestServer(fd, handlers)
+ if *testAllocator {
+ options := []RequestServerOption{WithRSAllocator()}
+ server, _ = NewRequestServerWithOptions(fd, handlers, options...)
+ } else {
+ server = NewRequestServer(fd, handlers)
+ }
  server.Serve()
  }()
  <-ready