Implement handler removal (#20)

Author: borzunov

conn.go

@@ -80,6 +80,14 @@ func (d *Daemon) handleConn(c net.Conn) {
 				return
 			}
 
+		case pb.Request_REMOVE_STREAM_HANDLER:
+			res := d.doRemoveStreamHandler(&req)
+			err := w.WriteMsg(res)
+			if err != nil {
+				log.Debugw("error writing response", "error", err)
+				return
+			}
+
 		case pb.Request_DHT:
 			res, ch, cancel := d.doDHT(&req)
 			err := w.WriteMsg(res)
@@ -271,22 +279,55 @@ func (d *Daemon) doStreamHandler(req *pb.Request) *pb.Response {
 	}
 	for _, sp := range req.StreamHandler.Proto {
 		p := protocol.ID(sp)
-		_, ok := d.handlers[p]
+		round_robin, ok := d.handlers[p]
 		if !ok {
 			d.handlers[p] = utils.NewRoundRobin()
-			d.handlers[p].Push(maddr)
+			d.handlers[p].Append(maddr)
 			d.host.SetStreamHandler(p, d.handleStream)
 		} else if !req.StreamHandler.GetBalanced() {
 			return errorResponseString(fmt.Sprintf("handler for protocol %s already set", p))
 		} else {
-			d.handlers[p].Push(maddr)
+			round_robin.Append(maddr)
 		}
 		log.Debugw("set stream handler", "protocol", sp, "to", maddr)
 	}
 
 	return okResponse()
 }
 
+func (d *Daemon) doRemoveStreamHandler(req *pb.Request) *pb.Response {
+	if req.RemoveStreamHandler == nil {
+		return errorResponseString("Malformed request; missing parameters")
+	}
+
+	d.mx.Lock()
+	defer d.mx.Unlock()
+
+	maddr, err := ma.NewMultiaddrBytes(req.RemoveStreamHandler.Addr)
+	if err != nil {
+		return errorResponse(err)
+	}
+	for _, sp := range req.RemoveStreamHandler.Proto {
+		p := protocol.ID(sp)
+		round_robin, ok := d.handlers[p]
+		if !ok {
+			return errorResponseString(fmt.Sprintf("handler for protocol %s does not exist", p))
+		}
+
+		ok = round_robin.Remove(maddr)
+		if !ok {
+			return errorResponseString(fmt.Sprintf("handler for protocol %s with maddr %s does not exist", p, maddr.String()))
+		}
+		if round_robin.Len() == 0 {
+			d.host.RemoveStreamHandler(p)
+			delete(d.handlers, p)
+		}
+		log.Debugw("removed stream handler", "protocol", sp, "to", maddr)
+	}
+
+	return okResponse()
+}
+
 func (d *Daemon) doListPeers(req *pb.Request) *pb.Response {
 	conns := d.host.Network().Conns()
 	peers := make([]*pb.PeerInfo, len(conns))

internal/utils/round_robin.go

@@ -1,5 +1,9 @@
 package utils
 
+import (
+	"reflect"
+)
+
 type RoundRobin struct {
 	data []interface{}
 	next int
@@ -11,10 +15,32 @@ func (r *RoundRobin) Next() interface{} {
 	return v
 }
 
-func (r *RoundRobin) Push(v interface{}) {
+func (r *RoundRobin) Append(v interface{}) {
 	r.data = append(r.data, v)
 }
 
+func (r *RoundRobin) Remove(v interface{}) bool {
+	found := -1
+	for index, item := range r.data {
+		if reflect.DeepEqual(item, v) {
+			found = index
+			break
+		}
+	}
+	if found == -1 {
+		return false
+	}
+
+	r.data = append(r.data[:found], r.data[found+1:]...)
+	if found < r.next {
+		r.next--
+	}
+	if r.next == r.Len() {
+		r.next = 0
+	}
+	return true
+}
+
 func (r *RoundRobin) Len() int {
 	return len(r.data)
 }

p2pclient/streams.go

@@ -222,3 +222,42 @@ func (c *Client) NewStreamHandler(protos []string, handler StreamHandlerFunc, ba
 
 	return nil
 }
+
+func (c *Client) RemoveStreamHandler(protos []string) error {
+	raw_control, err := c.newControlConn()
+	if err != nil {
+		return err
+	}
+	control := &byteReaderConn{raw_control}
+	defer control.Close()
+
+	c.mhandlers.Lock()
+	defer c.mhandlers.Unlock()
+
+	w := ggio.NewDelimitedWriter(control)
+	req := &pb.Request{
+		Type: pb.Request_REMOVE_STREAM_HANDLER.Enum(),
+		RemoveStreamHandler: &pb.RemoveStreamHandlerRequest{
+			Addr:  c.listenMaddr.Bytes(),
+			Proto: protos,
+		},
+	}
+	if err := w.WriteMsg(req); err != nil {
+		return err
+	}
+
+	resp := &pb.Response{}
+	err = readMsgSafe(control, resp)
+	if err != nil {
+		return err
+	}
+	if err := resp.GetError(); err != nil {
+		return fmt.Errorf("error from daemon: %s", err.GetMsg())
+	}
+
+	for _, proto := range protos {
+		delete(c.handlers, proto)
+	}
+
+	return nil
+}

p2pclient/unary_handlers.go

@@ -153,6 +153,31 @@ func (c *Client) AddUnaryHandler(proto protocol.ID, handler UnaryHandlerFunc, ba
 	return nil
 }
 
+func (c *Client) RemoveUnaryHandler(proto protocol.ID) error {
+	w := c.getPersistentWriter()
+
+	callID := uuid.New()
+
+	w.WriteMsg(
+		&pb.PersistentConnectionRequest{
+			CallId: callID[:],
+			Message: &pb.PersistentConnectionRequest_RemoveUnaryHandler{
+				RemoveUnaryHandler: &pb.RemoveUnaryHandlerRequest{
+					Proto: (*string)(&proto),
+				},
+			},
+		},
+	)
+
+	if _, err := c.getResponse(callID); err != nil {
+		return err
+	}
+
+	c.unaryHandlers.Delete(proto)
+
+	return nil
+}
+
 func (c *Client) CallUnaryHandler(
 	ctx context.Context,
 	peerID peer.ID,