Merge pull request #223 from yywing/issue_222

Closes 222

allocator.go

@@ -18,10 +18,10 @@ const (
 
 // allocator maintains a bitset of allocated numbers.
 type allocator struct {
- pool *big.Int
- last int
- low int
- high int
+ pool  *big.Int
+ follow int
+ low  int
+ high  int
 }
 
 // NewAllocator reserves and frees integers out of a range between low and
@@ -31,10 +31,10 @@ type allocator struct {
 // sizeof(big.Word)
 func newAllocator(low, high int) *allocator {
  return &allocator{
- pool: big.NewInt(0),
- last: low,
- low: low,
- high: high,
+ pool:  big.NewInt(0),
+ follow: low,
+ low:  low,
+ high:  high,
  }
 }
 
@@ -69,21 +69,29 @@ func (a allocator) String() string {
 // O(N) worst case runtime where N is allocated, but usually O(1) due to a
 // rolling index into the oldest allocation.
 func (a *allocator) next() (int, bool) {
- wrapped := a.last
+ wrapped := a.follow
+ defer func() {
+ // make a.follow point to next value
+ if a.follow == a.high {
+ a.follow = a.low
+ } else {
+ a.follow += 1
+ }
+ }()
 
  // Find trailing bit
- for ; a.last <= a.high; a.last++ {
- if a.reserve(a.last) {
- return a.last, true
+ for ; a.follow <= a.high; a.follow++ {
+ if a.reserve(a.follow) {
+ return a.follow, true
  }
  }
 
  // Find preceding free'd pool
- a.last = a.low
+ a.follow = a.low
 
- for ; a.last < wrapped; a.last++ {
- if a.reserve(a.last) {
- return a.last, true
+ for ; a.follow < wrapped; a.follow++ {
+ if a.reserve(a.follow) {
+ return a.follow, true
  }
  }
 

allocator_test.go

@@ -76,6 +76,28 @@ func TestAllocatorShouldReuseReleased(t *testing.T) {
  }
 }
 
+func TestAllocatorShouldNotReuseEarly(t *testing.T) {
+ a := newAllocator(1, 2)
+
+ first, _ := a.next()
+ if want, got := 1, first; want != got {
+ t.Fatalf("expected allocation to be %d, got: %d", want, got)
+ }
+
+ a.release(first)
+
+ second, _ := a.next()
+ if want, got := 2, second; want != got {
+ t.Fatalf("expected second allocation to be %d, got: %d", want, got)
+ }
+
+ third, _ := a.next()
+ if want, got := first, third; want != got {
+ t.Fatalf("expected third allocation to be %d, got: %d", want, got)
+ }
+
+}
+
 func TestAllocatorReleasesKeepUpWithAllocationsForAllSizes(t *testing.T) {
  if testing.Short() {
  t.Skip()

channel.go

@@ -41,6 +41,7 @@ type Channel struct {
 
  // closed is set to 1 when the channel has been closed - see Channel.send()
  closed int32
+ close chan struct{}
 
  // true when we will never notify again
  noNotify bool
@@ -86,6 +87,7 @@ func newChannel(c *Connection, id uint16) *Channel {
  confirms: newConfirms(),
  recv: (*Channel).recvMethod,
  errors: make(chan *Error, 1),
+ close: make(chan struct{}),
  }
 }
 
@@ -146,6 +148,7 @@ func (ch *Channel) shutdown(e *Error) {
  }
 
  close(ch.errors)
+ close(ch.close)
  ch.noNotify = true
  })
 }
@@ -368,7 +371,11 @@ func (ch *Channel) dispatch(msg message) {
  // deliveries are in flight and a no-wait cancel has happened
 
  default:
- ch.rpc <- msg
+ select {
+ case <-ch.close:
+ return
+ case ch.rpc <- msg:
+ }
  }
 }
 
@@ -468,6 +475,10 @@ code set to '200'.
 It is safe to call this method multiple times.
 */
 func (ch *Channel) Close() error {
+ if ch.IsClosed() {
+ return nil
+ }
+
  defer ch.connection.closeChannel(ch, nil)
  return ch.call(
  &channelClose{ReplyCode: replySuccess},

connection.go

@@ -112,6 +112,8 @@ type Connection struct {
  blocks []chan Blocking
 
  errors chan *Error
+ // if connection is closed should close this chan
+ close chan struct{}
 
  Config Config // The negotiated Config after connection.open
 
@@ -263,6 +265,7 @@ func Open(conn io.ReadWriteCloser, config Config) (*Connection, error) {
  rpc: make(chan message),
  sends: make(chan time.Time),
  errors: make(chan *Error, 1),
+ close: make(chan struct{}),
  deadlines: make(chan readDeadliner, 1),
  }
  go c.reader(conn)
@@ -597,6 +600,8 @@ func (c *Connection) shutdown(err *Error) {
  }
 
  c.conn.Close()
+ // reader exit
+ close(c.close)
 
  c.channels = nil
  c.allocator = nil
@@ -634,15 +639,23 @@ func (c *Connection) dispatch0(f frame) {
  c <- Blocking{Active: false}
  }
  default:
- c.rpc <- m
+ select {
+ case <-c.close:
+ return
+ case c.rpc <- m:
+ }
+
  }
  case *heartbeatFrame:
  // kthx - all reads reset our deadline. so we can drop this
  default:
  // lolwat - channel0 only responds to methods and heartbeats
- if err := c.closeWith(ErrUnexpectedFrame); err != nil {
- Logger.Printf("error sending connectionCloseOk with ErrUnexpectedFrame, error: %+v", err)
- }
+ // closeWith use call don't block reader
+ go func() {
+ if err := c.closeWith(ErrUnexpectedFrame); err != nil {
+ Logger.Printf("error sending connectionCloseOk with ErrUnexpectedFrame, error: %+v", err)
+ }
+ }()
  }
 }
 
@@ -689,9 +702,12 @@ func (c *Connection) dispatchClosed(f frame) {
  // we are already closed, so do nothing
  default:
  // unexpected method on closed channel
- if err := c.closeWith(ErrClosed); err != nil {
- Logger.Printf("error sending connectionCloseOk with ErrClosed, error: %+v", err)
- }
+ // closeWith use call don't block reader
+ go func() {
+ if err := c.closeWith(ErrClosed); err != nil {
+ Logger.Printf("error sending connectionCloseOk with ErrClosed, error: %+v", err)
+ }
+ }()
  }
  }
 }
@@ -813,13 +829,16 @@ func (c *Connection) allocateChannel() (*Channel, error) {
 
 // releaseChannel removes a channel from the registry as the final part of the
 // channel lifecycle
-func (c *Connection) releaseChannel(id uint16) {
+func (c *Connection) releaseChannel(ch *Channel) {
  c.m.Lock()
  defer c.m.Unlock()
 
  if !c.IsClosed() {
- delete(c.channels, id)
- c.allocator.release(int(id))
+ got, ok := c.channels[ch.id]
+ if ok && got == ch {
+ delete(c.channels, ch.id)
+ c.allocator.release(int(ch.id))
+ }
  }
 }
 
@@ -831,7 +850,7 @@ func (c *Connection) openChannel() (*Channel, error) {
  }
 
  if err := ch.open(); err != nil {
- c.releaseChannel(ch.id)
+ c.releaseChannel(ch)
  return nil, err
  }
  return ch, nil
@@ -842,7 +861,7 @@ func (c *Connection) openChannel() (*Channel, error) {
 // this connection.
 func (c *Connection) closeChannel(ch *Channel, e *Error) {
  ch.shutdown(e)
- c.releaseChannel(ch.id)
+ c.releaseChannel(ch)
 }
 
 /*
@@ -863,13 +882,14 @@ func (c *Connection) call(req message, res ...message) error {
  }
  }
 
- msg, ok := <-c.rpc
- if !ok {
-  err, errorsChanIsOpen := <-c.errors
- if !errorsChanIsOpen {
- return ErrClosed
+ var msg message
+ select {
+ case e, ok := <-c.errors:
+ if ok {
+ return e
  }
- return err
+ return ErrClosed
+ case msg = <-c.rpc:
  }
 
  // Try to match one of the result types