routing: support multiple namespaced MissionControls

routing/missioncontrol.go

@@ -135,12 +135,12 @@ type MissionControl struct {
 }
 
 // MissionController manages MissionControl instances in various namespaces.
-//
-// NOTE: currently it only has a MissionControl in the default namespace.
 type MissionController struct {
-	cfg *mcConfig
+	db           kvdb.Backend
+	cfg          *mcConfig
+	defaultMCCfg *MissionControlConfig
 
-	mc *MissionControl
+	mc map[string]*MissionControl
 	mu sync.Mutex
 
 	// TODO(roasbeef): further counters, if vertex continually unavailable,
@@ -149,12 +149,33 @@ type MissionController struct {
 	// TODO(roasbeef): also add favorable metrics for nodes
 }
 
-// GetDefaultStore returns the MissionControl in the default namespace.
-func (m *MissionController) GetDefaultStore() *MissionControl {
+// GetNamespacedStore returns the MissionControl in the given namespace. If one
+// does not yet exist, then it is initialised.
+func (m *MissionController) GetNamespacedStore(ns string) (*MissionControl,
+	error) {
+
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	return m.mc
+	if mc, ok := m.mc[ns]; ok {
+		return mc, nil
+	}
+
+	return m.initMissionControl(ns)
+}
+
+// ListNamespaces returns a list of the namespaces that the MissionController
+// is aware of.
+func (m *MissionController) ListNamespaces() []string {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	namespaces := make([]string, 0, len(m.mc))
+	for ns := range m.mc {
+		namespaces = append(namespaces, ns)
+	}
+
+	return namespaces
 }
 
 // MissionControlConfig defines parameters that control mission control
@@ -259,61 +280,143 @@ func NewMissionController(db kvdb.Backend, self route.Vertex,
 		return nil, err
 	}
 
+	mcCfg := &mcConfig{
+		clock:    clock.NewDefaultClock(),
+		selfNode: self,
+	}
+
+	mgr := &MissionController{
+		db:           db,
+		defaultMCCfg: cfg,
+		cfg:          mcCfg,
+		mc:           make(map[string]*MissionControl),
+	}
+
+	if err := mgr.loadMissionControls(); err != nil {
+		return nil, err
+	}
+
+	for _, mc := range mgr.mc {
+		if err := mc.init(); err != nil {
+			return nil, err
+		}
+	}
+
+	return mgr, nil
+}
+
+// loadMissionControls initialises a MissionControl in the default namespace if
+// one does not yet exist. It then initialises a MissionControl for all other
+// namespaces found in the DB.
+//
+// NOTE: this should only be called once during MissionController construction.
+func (m *MissionController) loadMissionControls() error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	// Always initialise the default namespace.
+	_, err := m.initMissionControl(DefaultMissionControlNamespace)
+	if err != nil {
+		return err
+	}
+
+	namespaces := make(map[string]struct{})
+	err = m.db.View(func(tx walletdb.ReadTx) error {
+		mcStoreBkt := tx.ReadBucket(resultsKey)
+		if mcStoreBkt == nil {
+			return fmt.Errorf("top level mission control bucket " +
+				"not found")
+		}
+
+		// Iterate through all the keys in the bucket and collect the
+		// namespaces.
+		return mcStoreBkt.ForEach(func(k, _ []byte) error {
+			// We've already initialised the default namespace so
+			// we can skip it.
+			if string(k) == DefaultMissionControlNamespace {
+				return nil
+			}
+
+			namespaces[string(k)] = struct{}{}
+
+			return nil
+		})
+	}, func() {})
+	if err != nil {
+		return err
+	}
+
+	// Now, iterate through all the namespaces and initialise them.
+	for ns := range namespaces {
+		_, err = m.initMissionControl(ns)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+// initMissionControl creates a new MissionControl instance with the given
+// namespace if one does not yet exist.
+//
+// NOTE: the MissionController's mutex must be held before calling this method.
+func (m *MissionController) initMissionControl(namespace string) (
+	*MissionControl, error) {
+
+	// If a mission control with this namespace has already been initialised
+	// then there is nothing left to do.
+	if mc, ok := m.mc[namespace]; ok {
+		return mc, nil
+	}
+
+	cfg := m.defaultMCCfg
+
 	store, err := newMissionControlStore(
-		newDefaultNamespacedStore(db), cfg.MaxMcHistory,
+		newNamespacedDB(m.db, namespace), cfg.MaxMcHistory,
 		cfg.McFlushInterval,
 	)
 	if err != nil {
 		return nil, err
 	}
 
-	mcCfg := &mcConfig{
-		clock:    clock.NewDefaultClock(),
-		selfNode: self,
-	}
-
-	// Create a mission control in the default namespace.
-	defaultMC := &MissionControl{
-		cfg:       mcCfg,
+	mc := &MissionControl{
+		cfg:       m.cfg,
 		state:     newMissionControlState(cfg.MinFailureRelaxInterval),
 		store:     store,
 		estimator: cfg.Estimator,
 		log: build.NewPrefixLog(
-			fmt.Sprintf("[%s]:", DefaultMissionControlNamespace),
-			log,
+			fmt.Sprintf("[%s]:", namespace), log,
 		),
 		onConfigUpdate: cfg.OnConfigUpdate,
 	}
 
-	mc := &MissionController{
-		cfg: mcCfg,
-		mc:  defaultMC,
-	}
-
-	if err := mc.mc.init(); err != nil {
-		return nil, err
-	}
+	m.mc[namespace] = mc
 
 	return mc, nil
 }
 
-// RunStoreTicker runs the mission control store's ticker.
-func (m *MissionController) RunStoreTicker() {
+// RunStoreTickers runs the mission controller store's tickers.
+func (m *MissionController) RunStoreTickers() {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	m.mc.store.run()
+	for _, mc := range m.mc {
+		mc.store.run()
+	}
 }
 
-// StopStoreTicker stops the mission control store's ticker.
-func (m *MissionController) StopStoreTicker() {
+// StopStoreTickers stops the mission control store's tickers.
+func (m *MissionController) StopStoreTickers() {
 	log.Debug("Stopping mission control store ticker")
 	defer log.Debug("Mission control store ticker stopped")
 
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
-	m.mc.store.stop()
+	for _, mc := range m.mc {
+		mc.store.stop()
+	}
 }
 
 // init initializes mission control with historical data.