Restructure version.Update to use channel

Author: mohamed-essam

client/internal/updatemanager/manager.go

@@ -35,7 +35,8 @@ type UpdateManager struct {
 	lastTrigger    time.Time
 	statusRecorder *peer.Status
 	mutex          sync.Mutex
-	waitGroup      sync.WaitGroup
+	updateChannel  chan string
+	doneChannel    chan struct{}
 }
 
 func NewUpdateManager(ctx context.Context, statusRecorder *peer.Status) *UpdateManager {
@@ -49,9 +50,12 @@ func NewUpdateManager(ctx context.Context, statusRecorder *peer.Status) *UpdateM
 		cancel:         cancel,
 		version:        disableAutoUpdate,
 		latestVersion:  unknownVersion,
+		updateChannel:  make(chan string, 4),
+		doneChannel:    make(chan struct{}),
 	}
 	update.SetDaemonVersion(version.NetbirdVersion())
-	update.SetOnUpdateListener(manager.Updated)
+	update.SetOnUpdateChannel(manager.updateChannel)
+	go manager.UpdateLoop()
 	return manager
 }
 
@@ -61,7 +65,7 @@ func (u *UpdateManager) SetVersion(v string) {
 		log.Tracef("Auto-update version set to %s", v)
 		u.version = v
 		u.mutex.Unlock()
-		go u.Updated(unknownVersion)
+		u.updateChannel <- unknownVersion
 	} else {
 		u.mutex.Unlock()
 	}
@@ -75,20 +79,26 @@ func (u *UpdateManager) Stop() {
 		u.update.StopWatch()
 		u.update = nil
 	}
-	u.waitGroup.Wait()
+	<-u.doneChannel
 }
 
-func (u *UpdateManager) Updated(latestVersion string) {
-	u.waitGroup.Add(1)
-	defer u.waitGroup.Done()
-	u.mutex.Lock()
-	defer u.mutex.Unlock()
-	if latestVersion != unknownVersion {
-		u.latestVersion = latestVersion
+func (u *UpdateManager) UpdateLoop() {
+	for {
+		select {
+		case <-u.ctx.Done():
+			u.doneChannel <- struct{}{}
+			return
+		case latestVersion := <-u.updateChannel:
+			u.mutex.Lock()
+			if latestVersion != unknownVersion {
+				u.latestVersion = latestVersion
+			}
+			u.mutex.Unlock()
+			ctx, cancel := context.WithDeadline(u.ctx, time.Now().Add(time.Minute))
+			u.CheckForUpdates(ctx)
+			cancel()
+		}
 	}
-	ctx, cancel := context.WithDeadline(u.ctx, time.Now().Add(time.Minute))
-	defer cancel()
-	u.CheckForUpdates(ctx)
 }
 
 func (u *UpdateManager) CheckForUpdates(ctx context.Context) {

client/ui/client_ui.go

@@ -1213,7 +1213,7 @@ func protoConfigToConfig(cfg *proto.GetConfigResponse) *profilemanager.Config {
 	return &config
 }
 
-func (s *serviceClient) onUpdateAvailable(_ string) {
+func (s *serviceClient) onUpdateAvailable() {
 	s.updateIndicationLock.Lock()
 	defer s.updateIndicationLock.Unlock()
 

management/internals/server/server.go

@@ -182,8 +182,8 @@ func (s *BaseServer) Start(ctx context.Context) error {
 
 	s.update = version.NewUpdate("nb/management")
 	s.update.SetDaemonVersion(version.NetbirdVersion())
-	s.update.SetOnUpdateListener(func(newVersion string) {
-		log.WithContext(ctx).Infof("your management version, \"%s\", is outdated, a new management version (%s) is available. Learn more here: https://github.com/netbirdio/netbird/releases", version.NetbirdVersion(), newVersion)
+	s.update.SetOnUpdateListener(func() {
+		log.WithContext(ctx).Infof("your management version, \"%s\", is outdated, a new management version is available. Learn more here: https://github.com/netbirdio/netbird/releases", version.NetbirdVersion())
 	})
 
 	return nil

version/update.go

@@ -30,7 +30,8 @@ type Update struct {
 	fetchTicker *time.Ticker
 	fetchDone   chan struct{}
 
-	onUpdateListener func(latestVersion string)
+	onUpdateListener func()
+	onUpdateChannel  chan string
 	listenerLock     sync.Mutex
 }
 
@@ -57,7 +58,11 @@ func NewUpdate(httpAgent string) *Update {
 // StopWatch stop the version info fetch loop
 func (u *Update) StopWatch() {
 	u.fetchTicker.Stop()
-	u.fetchDone <- struct{}{}
+
+	select {
+	case u.fetchDone <- struct{}{}:
+	default:
+	}
 }
 
 // SetDaemonVersion update the currently running daemon version. If new version is available it will trigger
@@ -80,15 +85,24 @@ func (u *Update) SetDaemonVersion(newVersion string) bool {
 }
 
 // SetOnUpdateListener set new update listener
-func (u *Update) SetOnUpdateListener(updateFn func(string)) {
+func (u *Update) SetOnUpdateListener(updateFn func()) {
 	u.listenerLock.Lock()
 	defer u.listenerLock.Unlock()
 
 	u.onUpdateListener = updateFn
+	if u.isUpdateAvailable() {
+		u.onUpdateListener()
+	}
+}
+
+func (u *Update) SetOnUpdateChannel(updateChannel chan string) {
+	u.listenerLock.Lock()
+	defer u.listenerLock.Unlock()
+	u.onUpdateChannel = updateChannel
 	if u.isUpdateAvailable() {
 		u.versionsLock.Lock()
 		defer u.versionsLock.Unlock()
-		u.onUpdateListener(u.latestAvailable.String())
+		u.onUpdateChannel <- u.latestAvailable.String()
 	}
 }
 
@@ -167,11 +181,14 @@ func (u *Update) checkUpdate() bool {
 
 	u.listenerLock.Lock()
 	defer u.listenerLock.Unlock()
+	if u.onUpdateChannel != nil {
+		u.onUpdateChannel <- u.latestAvailable.String()
+	}
 	if u.onUpdateListener == nil {
 		return true
 	}
 
-	go u.onUpdateListener(u.latestAvailable.String())
+	go u.onUpdateListener()
 	return true
 }
 

version/update_test.go

@@ -25,7 +25,7 @@ func TestNewUpdate(t *testing.T) {
 	onUpdate := false
 	u := NewUpdate(httpAgent)
 	defer u.StopWatch()
-	u.SetOnUpdateListener(func(_ string) {
+	u.SetOnUpdateListener(func() {
 		onUpdate = true
 		wg.Done()
 	})
@@ -50,7 +50,7 @@ func TestDoNotUpdate(t *testing.T) {
 	onUpdate := false
 	u := NewUpdate(httpAgent)
 	defer u.StopWatch()
-	u.SetOnUpdateListener(func(_ string) {
+	u.SetOnUpdateListener(func() {
 		onUpdate = true
 		wg.Done()
 	})
@@ -75,7 +75,7 @@ func TestDaemonUpdate(t *testing.T) {
 	onUpdate := false
 	u := NewUpdate(httpAgent)
 	defer u.StopWatch()
-	u.SetOnUpdateListener(func(_ string) {
+	u.SetOnUpdateListener(func() {
 		onUpdate = true
 		wg.Done()
 	})