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

@@ -468,6 +468,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

@@ -640,9 +640,12 @@ func (c *Connection) dispatch0(f frame) {
 		// 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 +692,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 +819,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 +840,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 +851,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 +872,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