client: add UpdateGroups to HeartbeatSender and PIMServer

Make UpdateGroups apply group changes in-place instead of tearing down
and restarting the goroutine/connection.

HeartbeatSender: mutex-protected dsts field, goroutine reads under lock
each tick, UpdateGroups swaps dsts and sends an immediate heartbeat.

PIMServer: channel-based update via updateCh, goroutine computes
added/removed groups and sends targeted join/prune messages. Eliminates
the need for callers to create a new RawConner on every update.

Remove conn param from PIMWriter.UpdateGroups interface and simplify
the multicast service update path accordingly.

Author: snormore

client/doublezerod/internal/manager/http_test.go

@@ -449,6 +449,10 @@ func (m *MockPIMServer) Start(conn pim.RawConner, iface string, tunnelAddr net.I
 	return nil
 }
 
+func (m *MockPIMServer) UpdateGroups(groups []net.IP) error {
+	return nil
+}
+
 func (m *MockPIMServer) Close() error {
 	return nil
 }
@@ -459,6 +463,10 @@ func (m *MockHeartbeatSender) Start(iface string, srcIP net.IP, groups []net.IP,
 	return nil
 }
 
+func (m *MockHeartbeatSender) UpdateGroups(groups []net.IP) error {
+	return nil
+}
+
 func (m *MockHeartbeatSender) Close() error {
 	return nil
 }

client/doublezerod/internal/manager/reconciler_test.go

@@ -69,14 +69,16 @@ func (m *mockBgpServer) GetPeerStatus(net.IP) bgp.Session          { return bgp.
 type mockPIMServer struct{}
 
 func (m *mockPIMServer) Start(pim.RawConner, string, net.IP, []net.IP) error { return nil }
+func (m *mockPIMServer) UpdateGroups([]net.IP) error                         { return nil }
 func (m *mockPIMServer) Close() error                                        { return nil }
 
 type mockHeartbeatSender struct{}
 
 func (m *mockHeartbeatSender) Start(string, net.IP, []net.IP, int, time.Duration) error {
 	return nil
 }
-func (m *mockHeartbeatSender) Close() error { return nil }
+func (m *mockHeartbeatSender) UpdateGroups([]net.IP) error { return nil }
+func (m *mockHeartbeatSender) Close() error                { return nil }
 
 // --- test helpers ---
 

client/doublezerod/internal/multicast/heartbeat.go

@@ -39,6 +39,9 @@ type PacketConner interface {
 type HeartbeatSender struct {
 	done chan struct{}
 	wg   *sync.WaitGroup
+	mu   sync.Mutex
+	dsts []*net.UDPAddr
+	conn PacketConner
 }
 
 func NewHeartbeatSender() *HeartbeatSender {
@@ -62,6 +65,24 @@ func (h *HeartbeatSender) Start(iface string, srcIP net.IP, groups []net.IP, ttl
 	return h.startWithConn(p, intf, groups, ttl, interval)
 }
 
+// UpdateGroups applies a new set of multicast groups in-place without
+// restarting the goroutine or connection. It sends an immediate heartbeat
+// to the new destinations so callers don't have to wait for the next tick.
+func (h *HeartbeatSender) UpdateGroups(groups []net.IP) error {
+	dsts := make([]*net.UDPAddr, len(groups))
+	for i, group := range groups {
+		dsts[i] = &net.UDPAddr{IP: group, Port: HeartbeatPort}
+	}
+
+	h.mu.Lock()
+	h.dsts = dsts
+	conn := h.conn
+	h.mu.Unlock()
+
+	sendHeartbeats(conn, dsts)
+	return nil
+}
+
 // startWithConn is the internal start method that accepts a pre-built connection for testing.
 func (h *HeartbeatSender) startWithConn(p PacketConner, intf *net.Interface, groups []net.IP, ttl int, interval time.Duration) error {
 	if err := p.SetMulticastTTL(ttl); err != nil {
@@ -78,21 +99,32 @@ func (h *HeartbeatSender) startWithConn(p PacketConner, intf *net.Interface, gro
 		dsts[i] = &net.UDPAddr{IP: group, Port: HeartbeatPort}
 	}
 
+	h.mu.Lock()
+	h.dsts = dsts
+	h.conn = p
+	h.mu.Unlock()
+
 	h.wg = &sync.WaitGroup{}
 	h.wg.Add(1)
 	go func() {
 		defer p.Close()
 		defer h.wg.Done()
 
 		// Send immediately before starting ticker so we don't delay by the interval.
-		sendHeartbeats(p, dsts)
+		h.mu.Lock()
+		currentDsts := h.dsts
+		h.mu.Unlock()
+		sendHeartbeats(p, currentDsts)
 
 		ticker := time.NewTicker(interval)
 		defer ticker.Stop()
 		for {
 			select {
 			case <-ticker.C:
-				sendHeartbeats(p, dsts)
+				h.mu.Lock()
+				currentDsts := h.dsts
+				h.mu.Unlock()
+				sendHeartbeats(p, currentDsts)
 			case <-h.done:
 				return
 			}

client/doublezerod/internal/multicast/heartbeat_test.go

@@ -264,6 +264,64 @@ func TestHeartbeatSender_DoubleClose(t *testing.T) {
 	sender.Close()
 }
 
+func TestHeartbeatSender_UpdateGroups(t *testing.T) {
+	conn := newMockPacketConn()
+	sender := NewHeartbeatSender()
+
+	// Start with one group.
+	intf := &net.Interface{Index: 1, Name: "lo0"}
+	err := sender.startWithConn(conn, intf, []net.IP{net.IPv4(239, 0, 0, 1)}, 32, 10*time.Second)
+	if err != nil {
+		t.Fatalf("failed to start: %v", err)
+	}
+	// Drain the immediate send.
+	<-conn.writeCh
+
+	// Update to two groups on the same running sender/connection.
+	if err := sender.UpdateGroups([]net.IP{net.IPv4(239, 0, 0, 1), net.IPv4(239, 0, 0, 2)}); err != nil {
+		t.Fatalf("UpdateGroups failed: %v", err)
+	}
+
+	// UpdateGroups sends an immediate heartbeat to all new destinations.
+	for i := range 2 {
+		select {
+		case <-conn.writeCh:
+		case <-time.After(2 * time.Second):
+			t.Fatalf("timed out waiting for heartbeat %d after update", i)
+		}
+	}
+
+	writes := conn.getWrites()
+	// At least 3 writes: 1 initial + 2 from UpdateGroups immediate send.
+	if len(writes) < 3 {
+		t.Fatalf("expected at least 3 writes, got %d", len(writes))
+	}
+
+	// Verify the last 2 writes went to both groups (from UpdateGroups immediate send).
+	seen := map[string]bool{}
+	for _, w := range writes[len(writes)-2:] {
+		udpAddr := w.dst.(*net.UDPAddr)
+		seen[udpAddr.IP.String()] = true
+	}
+	if !seen["239.0.0.1"] {
+		t.Error("missing heartbeat to 239.0.0.1 after update")
+	}
+	if !seen["239.0.0.2"] {
+		t.Error("missing heartbeat to 239.0.0.2 after update")
+	}
+
+	sender.Close()
+
+	// Verify the connection was NOT closed during UpdateGroups (same conn reused).
+	// It should only be closed after sender.Close().
+	conn.mu.Lock()
+	closed := conn.closed
+	conn.mu.Unlock()
+	if !closed {
+		t.Error("connection should be closed after sender.Close()")
+	}
+}
+
 func TestHeartbeatSender_CloseStopsSending(t *testing.T) {
 	conn := newMockPacketConn()
 	sender := NewHeartbeatSender()

client/doublezerod/internal/pim/server.go

@@ -30,21 +30,27 @@ type RawConner interface {
 }
 
 type PIMServer struct {
-	iface  string
-	groups []net.IP
-	done   chan struct{}
-	conn   RawConner
-	wg     *sync.WaitGroup
+	iface      string
+	groups     []net.IP
+	done       chan struct{}
+	conn       RawConner
+	wg         *sync.WaitGroup
+	tunnelAddr net.IP
+	updateCh   chan []net.IP
 }
 
 func NewPIMServer() *PIMServer {
-	return &PIMServer{done: make(chan struct{})}
+	return &PIMServer{
+		done:     make(chan struct{}),
+		updateCh: make(chan []net.IP),
+	}
 }
 
 func (s *PIMServer) Start(conn RawConner, iface string, tunnelAddr net.IP, groups []net.IP) error {
 	s.iface = iface
 	s.groups = groups
 	s.conn = conn
+	s.tunnelAddr = tunnelAddr
 
 	intf, err := net.InterfaceByName(s.iface)
 	if err != nil {
@@ -58,61 +64,96 @@ func (s *PIMServer) Start(conn RawConner, iface string, tunnelAddr net.IP, group
 	s.wg.Add(1)
 	go func() {
 		defer s.conn.Close()
+		defer s.wg.Done()
+
 		// send before we start ticker so we don't delay provisioning time by ticker interval
-		helloMsgBuf, err := constructHelloMessage()
-		if err != nil {
-			slog.Error("failed to serialize PIM hello msg", "error", err)
-		}
-		err = sendMsg(helloMsgBuf, intf, s.conn)
-		if err != nil {
-			slog.Error("failed to send PIM hello msg", "error", err)
-		}
-		joinPruneMsgBuf, err := constructJoinPruneMessage(tunnelAddr, groups, RpAddress, nil, joinHoldtime)
-		if err != nil {
-			slog.Error("failed to serialize PIM join msg", "error", err)
-		}
-		err = sendMsg(joinPruneMsgBuf, intf, s.conn)
-		if err != nil {
-			slog.Error("failed to send PIM join msg", "error", err)
-		}
+		sendHelloAndJoin(intf, s.conn, tunnelAddr, s.groups)
 
 		ticker := time.NewTicker(time.Second * 30)
+		defer ticker.Stop()
 		for {
 			select {
 			case <-ticker.C:
-				helloMsgBuf, err := constructHelloMessage()
-				if err != nil {
-					slog.Error("failed to serialize PIM hello msg", "error", err)
-				}
-				err = sendMsg(helloMsgBuf, intf, s.conn)
-				if err != nil {
-					slog.Error("failed to send PIM hello msg", "error", err)
-				}
-				joinPruneMsgBuf, err := constructJoinPruneMessage(tunnelAddr, groups, RpAddress, nil, joinHoldtime)
-				if err != nil {
-					slog.Error("failed to serialize PIM join msg", "error", err)
+				sendHelloAndJoin(intf, s.conn, tunnelAddr, s.groups)
+			case newGroups := <-s.updateCh:
+				added, removed := ipDiff(s.groups, newGroups)
+				if len(removed) > 0 {
+					buf, err := constructJoinPruneMessage(tunnelAddr, removed, nil, RpAddress, pruneHoldtime)
+					if err != nil {
+						slog.Error("failed to serialize PIM prune msg for removed groups", "error", err)
+					} else if err := sendMsg(buf, intf, s.conn); err != nil {
+						slog.Error("failed to send PIM prune msg for removed groups", "error", err)
+					}
 				}
-				err = sendMsg(joinPruneMsgBuf, intf, s.conn)
-				if err != nil {
-					slog.Error("failed to send PIM join msg", "error", err)
+				if len(added) > 0 {
+					buf, err := constructJoinPruneMessage(tunnelAddr, added, RpAddress, nil, joinHoldtime)
+					if err != nil {
+						slog.Error("failed to serialize PIM join msg for added groups", "error", err)
+					} else if err := sendMsg(buf, intf, s.conn); err != nil {
+						slog.Error("failed to send PIM join msg for added groups", "error", err)
+					}
 				}
+				s.groups = newGroups
 			case <-s.done:
-				joinPruneMsgBuf, err := constructJoinPruneMessage(tunnelAddr, groups, nil, RpAddress, pruneHoldtime)
+				joinPruneMsgBuf, err := constructJoinPruneMessage(tunnelAddr, s.groups, nil, RpAddress, pruneHoldtime)
 				if err != nil {
 					slog.Error("failed to serialize PIM prune msg", "error", err)
-				}
-				err = sendMsg(joinPruneMsgBuf, intf, s.conn)
-				if err != nil {
+				} else if err := sendMsg(joinPruneMsgBuf, intf, s.conn); err != nil {
 					slog.Error("failed to send PIM prune msg", "error", err)
 				}
-				s.wg.Done()
 				return
 			}
 		}
 	}()
 	return nil
 }
 
+// UpdateGroups applies a new set of multicast groups in-place without
+// restarting the goroutine or connection. The goroutine computes the diff
+// and sends targeted join/prune messages for added/removed groups.
+func (s *PIMServer) UpdateGroups(groups []net.IP) error {
+	s.updateCh <- groups
+	return nil
+}
+
+func sendHelloAndJoin(intf *net.Interface, conn RawConner, tunnelAddr net.IP, groups []net.IP) {
+	helloMsgBuf, err := constructHelloMessage()
+	if err != nil {
+		slog.Error("failed to serialize PIM hello msg", "error", err)
+	} else if err := sendMsg(helloMsgBuf, intf, conn); err != nil {
+		slog.Error("failed to send PIM hello msg", "error", err)
+	}
+	joinPruneMsgBuf, err := constructJoinPruneMessage(tunnelAddr, groups, RpAddress, nil, joinHoldtime)
+	if err != nil {
+		slog.Error("failed to serialize PIM join msg", "error", err)
+	} else if err := sendMsg(joinPruneMsgBuf, intf, conn); err != nil {
+		slog.Error("failed to send PIM join msg", "error", err)
+	}
+}
+
+// ipDiff returns IPs that were added and removed when transitioning from old to new.
+func ipDiff(old, new []net.IP) (added, removed []net.IP) {
+	oldSet := make(map[string]bool, len(old))
+	for _, ip := range old {
+		oldSet[ip.String()] = true
+	}
+	newSet := make(map[string]bool, len(new))
+	for _, ip := range new {
+		newSet[ip.String()] = true
+	}
+	for _, ip := range new {
+		if !oldSet[ip.String()] {
+			added = append(added, ip)
+		}
+	}
+	for _, ip := range old {
+		if !newSet[ip.String()] {
+			removed = append(removed, ip)
+		}
+	}
+	return
+}
+
 func (s *PIMServer) Close() error {
 	s.done <- struct{}{}
 	s.wg.Wait()