withdraw prefixes from RIBIn when RIB gets unregistered

routingtable/adjRIBIn/adj_rib_in.go

@@ -186,5 +186,13 @@ func (a *AdjRIBIn) Register(client routingtable.RouteTableClient) {
 
 // Unregister unregisters a client
 func (a *AdjRIBIn) Unregister(client routingtable.RouteTableClient) {
-	a.clientManager.Unregister(client)
+	if !a.clientManager.Unregister(client) {
+		return
+	}
+
+	for _, r := range a.rt.Dump() {
+		for _, p := range r.Paths() {
+			client.RemovePath(r.Prefix(), p)
+		}
+	}
 }

routingtable/adjRIBIn/adj_rib_in_test.go

@@ -152,7 +152,10 @@ func TestAddPath(t *testing.T) {
 		}
 
 		if test.removePath != nil {
-			removePathParams := mc.GetRemovePathParams()
+			r := mc.Removed()
+			assert.Equalf(t, 1, len(r), "Test %q failed: Call to RemovePath did not propagate prefix", test.name)
+
+			removePathParams := r[0]
 			if removePathParams.Pfx != test.removePfx {
 				t.Errorf("Test %q failed: Call to RemovePath did not propagate prefix properly: Got: %s Want: %s", test.name, removePathParams.Pfx.String(), test.removePfx.String())
 			}
@@ -296,13 +299,19 @@ func TestRemovePath(t *testing.T) {
 		adjRIBIn.RemovePath(test.removePfx, test.removePath)
 
 		if test.wantPropagation {
-			removePathParams := mc.GetRemovePathParams()
+			r := mc.Removed()
+			assert.Equalf(t, 1, len(r), "Test %q failed: Call to RemovePath did not propagate prefix", test.name)
+
+			removePathParams := r[0]
 			if removePathParams.Pfx != test.removePfx {
 				t.Errorf("Test %q failed: Call to RemovePath did not propagate prefix properly: Got: %s Want: %s", test.name, removePathParams.Pfx.String(), test.removePfx.String())
 			}
 			assert.Equal(t, test.removePath, removePathParams.Path)
 		} else {
-			removePathParams := mc.GetRemovePathParams()
+			r := mc.Removed()
+			assert.Equalf(t, 1, len(r), "Test %q failed: Call to RemovePath did not propagate prefix", test.name)
+
+			removePathParams := r[0]
 			uninitialized := net.Prefix{}
 			if removePathParams.Pfx != uninitialized {
 				t.Errorf("Test %q failed: Call to RemovePath propagated unexpectedly", test.name)
@@ -312,3 +321,44 @@ func TestRemovePath(t *testing.T) {
 		assert.Equal(t, test.expected, adjRIBIn.rt.Dump())
 	}
 }
+
+func TestUnregister(t *testing.T) {
+	adjRIBIn := New(filter.NewAcceptAllFilter(), routingtable.NewContributingASNs(), 0, 0, false)
+	mc := routingtable.NewRTMockClient()
+	adjRIBIn.Register(mc)
+
+	pfxs := []net.Prefix{
+		net.NewPfx(net.IPv4FromOctets(10, 0, 0, 0), 16),
+		net.NewPfx(net.IPv4FromOctets(10, 0, 1, 0), 24),
+	}
+
+	paths := []*route.Path{
+		&route.Path{
+			BGPPath: &route.BGPPath{
+				NextHop: net.IPv4FromOctets(192, 168, 0, 0),
+			},
+		},
+		&route.Path{
+			BGPPath: &route.BGPPath{
+				NextHop: net.IPv4FromOctets(192, 168, 2, 1),
+			},
+		},
+		&route.Path{
+			BGPPath: &route.BGPPath{
+				NextHop: net.IPv4FromOctets(192, 168, 3, 1),
+			},
+		},
+	}
+
+	adjRIBIn.RT().AddPath(pfxs[0], paths[0])
+	adjRIBIn.RT().AddPath(pfxs[0], paths[1])
+	adjRIBIn.RT().AddPath(pfxs[1], paths[2])
+
+	adjRIBIn.Unregister(mc)
+
+	r := mc.Removed()
+	assert.Equalf(t, 3, len(r), "Should have removed 3 paths, but only removed %d", len(r))
+	assert.Equal(t, &routingtable.RemovePathParams{pfxs[0], paths[0]}, r[0], "Withdraw 1")
+	assert.Equal(t, &routingtable.RemovePathParams{pfxs[0], paths[1]}, r[1], "Withdraw 2")
+	assert.Equal(t, &routingtable.RemovePathParams{pfxs[1], paths[2]}, r[2], "Withdraw 3")
+}

routingtable/mock_client.go

@@ -13,15 +13,17 @@ type RemovePathParams struct {
 }
 
 type RTMockClient struct {
-	removePathParams RemovePathParams
+	removed []*RemovePathParams
 }
 
 func NewRTMockClient() *RTMockClient {
-	return &RTMockClient{}
+	return &RTMockClient{
+		removed: make([]*RemovePathParams, 0),
+	}
 }
 
-func (m *RTMockClient) GetRemovePathParams() RemovePathParams {
-	return m.removePathParams
+func (m *RTMockClient) Removed() []*RemovePathParams {
+	return m.removed
 }
 
 func (m *RTMockClient) AddPath(pfx net.Prefix, p *route.Path) error {
@@ -46,8 +48,12 @@ func (m *RTMockClient) Unregister(RouteTableClient) {
 
 // RemovePath removes the path for prefix `pfx`
 func (m *RTMockClient) RemovePath(pfx net.Prefix, p *route.Path) bool {
-	m.removePathParams.Pfx = pfx
-	m.removePathParams.Path = p
+	params := &RemovePathParams{
+		Path: p,
+		Pfx:  pfx,
+	}
+	m.removed = append(m.removed, params)
+
 	return true
 }