From 279fba0c5741421b77e0798c0dd7485cbc110c99 Mon Sep 17 00:00:00 2001
From: ston1th <ston1th@giftfish.de>
Date: Sun, 26 Sep 2021 21:31:05 +0200
Subject: [PATCH] shutdown fix

---
 TODO                           |  2 ++
 cmd/haproxy-lb/main.go         |  4 ++--
 pkg/controller/loadbalancer.go |  7 +------
 pkg/vip/vip.go                 | 37 +++++++++++++++++++++++++++++++---
 4 files changed, 39 insertions(+), 11 deletions(-)

diff --git a/TODO b/TODO
index c7b2220..19a472d 100644
--- a/TODO
+++ b/TODO
@@ -2,4 +2,6 @@ VIP for haproxy-lb
 * take from pool?
 * extra config?
 
+cleanup pkg/vip
+
 test api in followers
diff --git a/cmd/haproxy-lb/main.go b/cmd/haproxy-lb/main.go
index 2fd6877..1f3ea33 100644
--- a/cmd/haproxy-lb/main.go
+++ b/cmd/haproxy-lb/main.go
@@ -84,8 +84,8 @@ func main() {
 	log.Info("haproxy-lb shutdown", "signal", s.String())
 	go func() {
 		select {
-		case <-time.After(time.Second * 10):
-			log.Error(errors.New("shutdown took longer than 10 seconds"), "shutdown forced")
+		case <-time.After(time.Second * 20):
+			log.Error(errors.New("shutdown took longer than 20 seconds"), "shutdown forced")
 		case <-sigs:
 			log.Error(errors.New("force shutdown triggerd"), "shutdown forced")
 		}
diff --git a/pkg/controller/loadbalancer.go b/pkg/controller/loadbalancer.go
index b591833..18328ce 100644
--- a/pkg/controller/loadbalancer.go
+++ b/pkg/controller/loadbalancer.go
@@ -129,12 +129,7 @@ func NewLBController(cfg *config.Config, srv *api.Server, log logr.Logger) (call
 			}
 		},
 		Cleanup: func(ctx context.Context, cc cluster.CallbackContext) {
-			db := db.New(cc)
-			ips, err := db.GetIPs()
-			if err != nil {
-				return
-			}
-			deleteIPs(cc, n, ips)
+			deleteIPs(cc, n, n.GetIPs())
 		},
 	}
 	return
diff --git a/pkg/vip/vip.go b/pkg/vip/vip.go
index ee1a639..6090d9f 100644
--- a/pkg/vip/vip.go
+++ b/pkg/vip/vip.go
@@ -2,8 +2,10 @@ package vip
 
 import (
 	"errors"
-	"github.com/vishvananda/netlink"
 	"net"
+	"sync"
+
+	"github.com/vishvananda/netlink"
 )
 
 var ErrNoDefaultInterface = errors.New("no default interface found")
@@ -25,6 +27,8 @@ func DefaultInterface() (iface string, err error) {
 }
 
 type Network struct {
+	sync.Mutex
+	m     map[string]struct{}
 	iface string
 	label string
 	link  netlink.Link
@@ -45,7 +49,12 @@ func NewNetworkWithLabel(iface, label string) (n *Network, err error) {
 	if err != nil {
 		return
 	}
-	n = &Network{iface, label, link}
+	n = &Network{
+		m:     make(map[string]struct{}),
+		iface: iface,
+		label: label,
+		link:  link,
+	}
 	return
 }
 
@@ -67,6 +76,8 @@ func (n *Network) HasIP(a *netlink.Addr) bool {
 }
 
 func (n *Network) AddIP(cidr string) error {
+	n.Lock()
+	defer n.Unlock()
 	addr, err := netlink.ParseAddr(cidr)
 	if err != nil {
 		return err
@@ -81,10 +92,13 @@ func (n *Network) AddIP(cidr string) error {
 	if err != nil {
 		return err
 	}
+	n.m[cidr] = struct{}{}
 	return ARPSendGratuitous(cidr, n.iface)
 }
 
 func (n *Network) DeleteIP(cidr string) error {
+	n.Lock()
+	defer n.Unlock()
 	addr, err := netlink.ParseAddr(cidr)
 	if err != nil {
 		return err
@@ -92,5 +106,22 @@ func (n *Network) DeleteIP(cidr string) error {
 	if !n.HasIP(addr) {
 		return nil
 	}
-	return netlink.AddrDel(n.link, addr)
+	err = netlink.AddrDel(n.link, addr)
+	if err != nil {
+		return err
+	}
+	delete(n.m, cidr)
+	return nil
+}
+
+func (n *Network) GetIPs() (cidrs []string) {
+	n.Lock()
+	defer n.Unlock()
+	cidrs = make([]string, len(n.m))
+	c := 0
+	for k := range n.m {
+		cidrs[c] = k
+		c++
+	}
+	return
 }