1 files changed, 78 insertions, 2 deletions
diff --git a/client/discovery.go b/client/discovery.go
index 415be9b..2758678 100644
--- a/client/discovery.go
+++ b/client/discovery.go
@@ -1,9 +1,13 @@
 package client
 
 import (
+	"context"
 	"sort"
 	"strings"
+	"sync"
 
+	"github.com/eduvpn/eduvpn-common/internal/discovery"
+	"github.com/eduvpn/eduvpn-common/internal/log"
 	"github.com/eduvpn/eduvpn-common/i18nerr"
 	"github.com/eduvpn/eduvpn-common/types/cookie"
 	discotypes "github.com/eduvpn/eduvpn-common/types/discovery"
@@ -24,7 +28,10 @@ func (c *Client) DiscoOrganizations(ck *cookie.Cookie, search string) (*discotyp
 		return nil, i18nerr.NewInternal("Server/organization discovery with this client ID is not supported")
 	}
 
-	orgs, fresh, err := c.cfg.Discovery().Organizations(ck.Context())
+	disco, release := c.discoMan.Discovery(true)
+	defer release()
+
+	orgs, fresh, err := disco.Organizations(ck.Context())
 	if fresh {
 		defer c.TrySave()
 	}
@@ -70,7 +77,9 @@ func (c *Client) DiscoServers(ck *cookie.Cookie, search string) (*discotypes.Ser
 		return nil, i18nerr.NewInternal("Server/organization discovery with this client ID is not supported")
 	}
 
-	servs, fresh, err := c.cfg.Discovery().Servers(ck.Context())
+	disco, release := c.discoMan.Discovery(true)
+	defer release()
+	servs, fresh, err := disco.Servers(ck.Context())
 	if fresh {
 		defer c.TrySave()
 	}
@@ -105,3 +114,70 @@ func (c *Client) DiscoServers(ck *cookie.Cookie, search string) (*discotypes.Ser
 		List: retServs,
 	}, err
 }
+
+type DiscoManager struct {
+	disco *discovery.Discovery
+
+	cancel context.CancelFunc
+	mu sync.RWMutex
+	wait sync.WaitGroup
+}
+
+func (m *DiscoManager) lock(write bool) {
+	if write {
+		m.mu.Lock()
+		return
+	}
+	m.mu.RLock()
+}
+
+func (m *DiscoManager) unlock(write bool) {
+	if write {
+		m.mu.Unlock()
+		return
+	}
+	m.mu.RUnlock()
+}
+
+func (m *DiscoManager) Discovery(write bool) (*discovery.Discovery, func()) {
+	if write {
+		m.wait.Wait()
+	}
+	m.lock(write)
+	return m.disco, func() {
+		m.unlock(write)
+	}
+}
+
+func (m *DiscoManager) Cancel() {
+	if m.cancel != nil {
+		m.cancel()
+	}
+	m.wait.Wait()
+}
+
+func (m *DiscoManager) Startup(ctx context.Context, cb func()) {
+	ctx, cancel := context.WithCancel(ctx)
+	m.cancel = cancel
+	m.wait.Add(1)
+	go func() {
+		defer m.wait.Done()
+		m.lock(false)
+		discoCopy, err := m.disco.Copy()
+		if err != nil {
+			log.Logger.Warningf("internal error, failed to clone discovery, %v", err)
+			return
+		}
+		m.unlock(false)
+		// we already log the warning
+		discoCopy.Servers(ctx) //nolint:errcheck
+
+		m.lock(true)
+		m.disco.UpdateServers(discoCopy)
+		m.unlock(true)
+
+		if cb != nil {
+			cb()
+		}
+	}()
+}