gotelem/mprpc/rpcconntrack.go

package mprpc

import (
	"errors"
	"math/rand"
	"sync"
)

// RPCConntrack is a request-response tracker that is used to connect
// the response to the appropriate caller.
type rpcConnTrack struct {
	ct map[uint32]chan Response
	mu sync.RWMutex
}


func NewRPCConnTrack() rpcConnTrack {
	return rpcConnTrack{
		ct: make(map[uint32]chan Response),
	}
}

// Get attempts to get a random mark from the mutex.
func (c *rpcConnTrack) Claim() (uint32, chan Response) {
	var val uint32
	for {

		//
		newVal := rand.Uint32()

		// BUG(saji): rpcConnTrack collisions are inefficient.

		// collision is *rare* - so we just try again.
		// I hope to god you don't saturate this tracker.
		c.mu.RLock()
		if _, exist := c.ct[newVal]; !exist {
			val = newVal
			c.mu.RUnlock()
			break
		}
		c.mu.RUnlock()
	}

	// claim it
	// the channel should be buffered. We only expect one value to go through.
	// so the size is fixed to 1.
	ch := make(chan Response, 1)
	c.mu.Lock()
	c.ct[val] = ch
	c.mu.Unlock()

	return val, ch
}

// Clear deletes the connection from the tracker and returns the channel
// associated with it. The caller can use the channel afterwards
// to send the response. It is the caller's responsibility to close the channel.
func (c *rpcConnTrack) Clear(val uint32) (chan Response, error) {
	c.mu.RLock()
	ch, ok := c.ct[val]
	c.mu.RUnlock()
	if !ok {
		return nil, errors.New("invalid msg id")
	}
	c.mu.Lock()
	delete(c.ct, val)
	c.mu.Unlock()
	return ch, nil
}