Simplify proxy poll handler, and broker match test

broker/broker.go

@@ -29,7 +29,7 @@ type BrokerContext struct {
 	// Map keeping track of snowflakeIDs required to match SDP answers from
 	// the second http POST.
 	snowflakeMap map[string]*Snowflake
-	createChannel chan *ProxyRequest
+	proxyPolls   chan *ProxyPoll
 }
 
 func NewBrokerContext() *BrokerContext {
@@ -38,7 +38,7 @@ func NewBrokerContext() *BrokerContext {
 	return &BrokerContext{
 		snowflakes:   snowflakes,
 		snowflakeMap: make(map[string]*Snowflake),
-		createChannel: make(chan *ProxyRequest),
+		proxyPolls:   make(chan *ProxyPoll),
 	}
 }
 
@@ -51,46 +51,58 @@ func (sh SnowflakeHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	sh.h(sh.BrokerContext, w, r)
 }
 
-type ProxyRequest struct {
+// Proxies may poll for client offers concurrently.
+type ProxyPoll struct {
 	id           string
-	offerChan chan []byte
+	offerChannel chan []byte
+}
+
+// Registers a Snowflake and waits for some Client to send an offer,
+// as part of the polling logic of the proxy handler.
+func (ctx *BrokerContext) RequestOffer(id string) []byte {
+	request := new(ProxyPoll)
+	request.id = id
+	request.offerChannel = make(chan []byte)
+	ctx.proxyPolls <- request
+	// Block until an offer is available...
+	offer := <-request.offerChannel
+	return offer
+}
+
+// goroutine which match proxies to clients.
+// Safely processes proxy requests, responding to them with either an available
+// client offer or nil on timeout / none are available.
+func (ctx *BrokerContext) Broker() {
+	for request := range ctx.proxyPolls {
+		snowflake := ctx.AddSnowflake(request.id)
+		// Wait for a client to avail an offer to the snowflake.
+		go func(request *ProxyPoll) {
+			select {
+			case offer := <-snowflake.offerChannel:
+				log.Println("Passing client offer to snowflake proxy.")
+				request.offerChannel <- offer
+			case <-time.After(time.Second * ProxyTimeout):
+				// This snowflake is no longer available to serve clients.
+				heap.Remove(ctx.snowflakes, snowflake.index)
+				delete(ctx.snowflakeMap, snowflake.id)
+				request.offerChannel <- nil
+			}
+		}(request)
+	}
 }
 
 // Create and add a Snowflake to the heap.
-func (sc *BrokerContext) AddSnowflake(id string) *Snowflake {
+func (ctx *BrokerContext) AddSnowflake(id string) *Snowflake {
 	snowflake := new(Snowflake)
 	snowflake.id = id
 	snowflake.clients = 0
 	snowflake.offerChannel = make(chan []byte)
 	snowflake.answerChannel = make(chan []byte)
-	heap.Push(sc.snowflakes, snowflake)
+	heap.Push(ctx.snowflakes, snowflake)
-	sc.snowflakeMap[id] = snowflake
+	ctx.snowflakeMap[id] = snowflake
 	return snowflake
 }
 
-// Match proxies to clients.
-// func (ctx *BrokerContext) Broker(proxies <-chan *ProxyRequest) {
-func (ctx *BrokerContext) Broker() {
-	// for p := range proxies {
-	for p := range ctx.createChannel {
-		snowflake := ctx.AddSnowflake(p.id)
-		// Wait for a client to avail an offer to the snowflake, or timeout
-		// and ask the snowflake to poll later.
-		go func(p *ProxyRequest) {
-			select {
-			case offer := <-snowflake.offerChannel:
-				log.Println("Passing client offer to snowflake.")
-				p.offerChan <- offer
-			case <-time.After(time.Second * ProxyTimeout):
-				// This snowflake is no longer available to serve clients.
-				heap.Remove(ctx.snowflakes, snowflake.index)
-				delete(ctx.snowflakeMap, snowflake.id)
-				p.offerChan <- nil
-			}
-		}(p)
-	}
-}
-
 func robotsTxtHandler(w http.ResponseWriter, r *http.Request) {
 	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
 	w.Write([]byte("User-agent: *\nDisallow:\n"))
@@ -145,14 +157,15 @@ func clientHandler(ctx *BrokerContext, w http.ResponseWriter, r *http.Request) {
 	case answer := <-snowflake.answerChannel:
 		log.Println("Client: Retrieving answer")
 		w.Write(answer)
-		// Only remove from the snowflake map once the answer is set.
-		delete(ctx.snowflakeMap, snowflake.id)
 
 	case <-time.After(time.Second * ClientTimeout):
 		log.Println("Client: Timed out.")
 		w.WriteHeader(http.StatusGatewayTimeout)
 		w.Write([]byte("timed out waiting for answer!"))
 	}
+	// Remove from the snowflake map whether answer was sent or not, because
+	// this client request is now over.
+	delete(ctx.snowflakeMap, snowflake.id)
 }
 
 /*
@@ -172,17 +185,9 @@ func proxyHandler(ctx *BrokerContext, w http.ResponseWriter, r *http.Request) {
 	if string(body) != id { // Mismatched IDs!
 		w.WriteHeader(http.StatusBadRequest)
 	}
-	// Maybe confirm that X-Session-ID is the same.
 	log.Println("Received snowflake: ", id)
-
+	// Wait for a client to avail an offer to the snowflake, or timeout if nil.
-	p := new(ProxyRequest)
+	offer := ctx.RequestOffer(id)
-	p.id = id
-	p.offerChan = make(chan []byte)
-	ctx.createChannel <- p
-
-	// Wait for a client to avail an offer to the snowflake, or timeout
-	// and ask the snowflake to poll later.
-	offer := <-p.offerChan
 	if nil == offer {
 		log.Println("Proxy " + id + " did not receive a Client offer.")
 		w.WriteHeader(http.StatusGatewayTimeout)

broker/snowflake-broker_test.go

@@ -22,6 +22,19 @@ func TestBroker(t *testing.T) {
 			So(len(ctx.snowflakeMap), ShouldEqual, 1)
 		})
 
+		Convey("Broker goroutine matches clients with proxies", func() {
+			p := new(ProxyPoll)
+			p.id = "test"
+			p.offerChannel = make(chan []byte)
+			go func() {
+				ctx.proxyPolls <- p
+				close(ctx.proxyPolls)
+			}()
+			ctx.Broker()
+			So(ctx.snowflakes.Len(), ShouldEqual, 1)
+			So(ctx.snowflakes.Len(), ShouldEqual, 1)
+		})
+
 		Convey("Responds to client offers...", func() {
 			w := httptest.NewRecorder()
 			data := bytes.NewReader([]byte("test"))
@@ -83,9 +96,9 @@ func TestBroker(t *testing.T) {
 					done <- true
 				}(ctx)
 				// Pass a fake client offer to this proxy
-				p := <-ctx.createChannel
+				p := <-ctx.proxyPolls
 				So(p.id, ShouldEqual, "test")
-				p.offerChan <- []byte("fake offer")
+				p.offerChannel <- []byte("fake offer")
 				<-done
 				So(w.Code, ShouldEqual, http.StatusOK)
 				So(w.Body.String(), ShouldEqual, "fake offer")
@@ -96,10 +109,10 @@ func TestBroker(t *testing.T) {
 					proxyHandler(ctx, w, r)
 					done <- true
 				}(ctx)
-				p := <-ctx.createChannel
+				p := <-ctx.proxyPolls
 				So(p.id, ShouldEqual, "test")
 				// nil means timeout
-				p.offerChan <- nil
+				p.offerChannel <- nil
 				<-done
 				So(w.Body.String(), ShouldEqual, "")
 				So(w.Code, ShouldEqual, http.StatusGatewayTimeout)
@@ -159,12 +172,12 @@ func TestBroker(t *testing.T) {
 		}()
 
 		// Manually do the Broker goroutine action here for full control.
-		p := <-ctx.createChannel
+		p := <-ctx.proxyPolls
 		So(p.id, ShouldEqual, "test")
 		s := ctx.AddSnowflake(p.id)
 		go func() {
 			offer := <-s.offerChannel
-			p.offerChan <- offer
+			p.offerChannel <- offer
 		}()
 		So(ctx.snowflakeMap["test"], ShouldNotBeNil)
 

broker/snowflake-heap.go

@@ -4,6 +4,10 @@ Keeping track of pending available snowflake proxies.
 
 package snowflake_broker
 
+/*
+The Snowflake struct contains a single interaction
+over the offer and answer channels.
+*/
 type Snowflake struct {
 	id            string
 	offerChannel  chan []byte