Safely write to response channels

config.go

@@ -6,7 +6,7 @@ import (
 
 const (
 	// With compression: burlesque.kch#opts=c#zcomp=gz#msiz=524288000
-	DefaultProductionStorage = "burlesque.kch#dfunit=100#msiz=512M"
+	DefaultProductionStorage = "burlesque.kch#dfunit=8#msiz=512M"
 )
 
 var (

request.go

@@ -6,10 +6,10 @@ import (
 
 type (
 	Request struct {
-		Queues   []string
-		Callback func(*Response)
-		Abort    chan bool
-		Dead     bool
+		Queues     []string
+		ResponseCh chan Response
+		Abort      chan bool
+		Dead       bool
 	}
 	Response struct {
 		Queue   string
@@ -25,37 +25,61 @@ var (
 )
 
 func RegisterPublication(q string, msg Message) bool {
-	pool.mutex.Lock()
-	for i, r := range pool.Requests {
-		for _, queueName := range r.Queues {
-			if queueName == q {
-				go r.Callback(&Response{Queue: queueName, Message: msg})
-				pool.Requests = append(pool.Requests[:i], pool.Requests[i+1:]...)
-				defer pool.mutex.Unlock()
-
-				return true
+	for _, r := range pool.Requests {
+		if r.Dead {
+			continue
+		}
+		for _, qname := range r.Queues {
+			if qname == q {
+				rsp := Response{Queue: q, Message: msg}
+				ok := r.TryRespond(rsp)
+				if ok {
+					return true
+				}
 			}
 		}
 	}
-	pool.mutex.Unlock()
 
 	ok := GetQueue(q).Push(msg)
 	return ok
 }
 
 func RegisterSubscription(r *Request) {
-	for _, queueName := range r.Queues {
-		q := GetQueue(queueName)
+	for _, qname := range r.Queues {
+		q := GetQueue(qname)
 		msg, ok := q.TryFetch(r.Abort)
 		if ok {
-			go r.Callback(&Response{Queue: queueName, Message: msg})
+			rsp := Response{Queue: qname, Message: msg}
+			ok := r.TryRespond(rsp)
+			if !ok {
+				q.Push(msg)
+			}
+
 			return
 		}
 	}
 
-	pool.mutex.Lock()
 	pool.Requests = append(pool.Requests, r)
-	pool.mutex.Unlock()
+}
+
+func (r *Request) TryRespond(rsp Response) bool {
+	okch := make(chan bool)
+
+	go func() {
+		defer func() {
+			err := recover()
+			if err != nil {
+				r.Dead = true
+				okch <- false
+			}
+		}()
+
+		r.ResponseCh <- rsp
+		okch <- true
+	}()
+
+	ok := <-okch
+	return ok
 }
 
 func (r *Request) Purge() {

server.go

@@ -49,27 +49,25 @@ func PublishHandler(w http.ResponseWriter, r *http.Request) {
 		msg = Message(r.FormValue("msg"))
 	}
 
-	queueName := r.FormValue("queue")
-	ok := RegisterPublication(queueName, msg)
+	qname := r.FormValue("queue")
+	ok := RegisterPublication(qname, msg)
 
 	if ok {
-		Debug("Published message of %d bytes to queue %s", len(msg), queueName)
+		Debug("Published message of %d bytes to queue %s", len(msg), qname)
 		w.Write([]byte("OK"))
 	} else {
-		Debug("Failed to publish message of %d bytes to queue %s", len(msg), queueName)
+		Debug("Failed to publish message of %d bytes to queue %s", len(msg), qname)
 		http.Error(w, "FAIL", 500)
 	}
 }
 
 func SubscriptionHandler(w http.ResponseWriter, r *http.Request) {
-	rch := make(chan *Response)
+	rch := make(chan Response)
 	abort := make(chan bool, 1)
 	req := &Request{
-		Queues: strings.Split(r.FormValue("queues"), ","),
-		Callback: func(r *Response) {
-			rch <- r
-		},
-		Abort: abort,
+		Queues:     strings.Split(r.FormValue("queues"), ","),
+		ResponseCh: rch,
+		Abort:      abort,
 	}
 	go RegisterSubscription(req)
 
@@ -78,15 +76,15 @@ func SubscriptionHandler(w http.ResponseWriter, r *http.Request) {
 	go func() {
 		select {
 		case <-disconnected:
-			rch <- nil
+			close(rch)
 			abort <- true
-			req.Purge()
 		case <-finished:
 		}
+		req.Purge()
 	}()
 
-	res := <-rch
-	if res == nil {
+	res, ok := <-rch
+	if !ok {
 		return
 	}
 

state.go

@@ -48,8 +48,8 @@ func LoadState() {
 		return
 	}
 
-	for queueName, meta := range state {
-		RegisterQueue(queueName, meta["wi"], meta["ri"])
+	for qname, meta := range state {
+		RegisterQueue(qname, meta["wi"], meta["ri"])
 	}
 
 	Log("State successfully loaded")