Implement UnsetForSource method (DataDog#31733)

pkg/config/model/viper.go

@@ -91,9 +91,18 @@ type ValueWithSource struct {
 	Value  interface{}
 }
 
-// IsGreaterOrEqualThan returns true if the current source is of higher priority than the one given as a parameter
-func (s Source) IsGreaterOrEqualThan(x Source) bool {
-	return sourcesPriority[s] >= sourcesPriority[x]
+// IsGreaterThan returns true if the current source is of higher priority than the one given as a parameter
+func (s Source) IsGreaterThan(x Source) bool {
+	return sourcesPriority[s] > sourcesPriority[x]
+}
+
+// PreviousSource returns the source before the current one, or Default (lowest priority) if there isn't one
+func (s Source) PreviousSource() Source {
+	previous := sourcesPriority[s]
+	if previous == 0 {
+		return sources[previous]
+	}
+	return sources[previous-1]
 }
 
 // String casts Source into a string

pkg/config/nodetreemodel/config.go

@@ -241,11 +241,83 @@ func (c *ntmConfig) SetDefault(key string, value interface{}) {
 	_, _ = c.defaults.SetAt(parts, value, model.SourceDefault)
 }
 
+func (c *ntmConfig) findPreviousSourceNode(key string, source model.Source) (Node, error) {
+	iter := source
+	for iter != model.SourceDefault {
+		iter = iter.PreviousSource()
+		tree, err := c.getTreeBySource(iter)
+		if err != nil {
+			return nil, err
+		}
+		node := c.leafAtPathFromNode(key, tree)
+		if _, isMissing := node.(*missingLeafImpl); !isMissing {
+			return node, nil
+		}
+	}
+	return nil, ErrNotFound
+}
+
 // UnsetForSource unsets a config entry for a given source
-func (c *ntmConfig) UnsetForSource(_key string, _source model.Source) {
+func (c *ntmConfig) UnsetForSource(key string, source model.Source) {
 	c.Lock()
-	c.logErrorNotImplemented("UnsetForSource")
-	c.Unlock()
+	defer c.Unlock()
+
+	// Remove it from the original source tree
+	tree, err := c.getTreeBySource(source)
+	if err != nil {
+		log.Errorf("%s", err)
+		return
+	}
+	parentNode, childName, err := c.parentOfNode(tree, key)
+	if err != nil {
+		return
+	}
+	// Only remove if the setting is a leaf
+	if child, err := parentNode.GetChild(childName); err == nil {
+		if _, ok := child.(LeafNode); ok {
+			parentNode.RemoveChild(childName)
+		} else {
+			log.Errorf("cannot remove setting %q, not a leaf", key)
+			return
+		}
+	}
+
+	// If the node in the merged tree doesn't match the source we expect, we're done
+	if c.leafAtPathFromNode(key, c.root).Source() != source {
+		return
+	}
+
+	// Find what the previous value used to be, based upon the previous source
+	prevNode, err := c.findPreviousSourceNode(key, source)
+	if err != nil {
+		return
+	}
+
+	// Get the parent node of the leaf we're unsetting
+	parentNode, childName, err = c.parentOfNode(c.root, key)
+	if err != nil {
+		return
+	}
+	// Replace the child with the node from the previous layer
+	parentNode.InsertChildNode(childName, prevNode.Clone())
+}
+
+func (c *ntmConfig) parentOfNode(node Node, key string) (InnerNode, string, error) {
+	parts := splitKey(key)
+	lastPart := parts[len(parts)-1]
+	parts = parts[:len(parts)-1]
+	var err error
+	for _, p := range parts {
+		node, err = node.GetChild(p)
+		if err != nil {
+			return nil, "", err
+		}
+	}
+	innerNode, ok := node.(InnerNode)
+	if !ok {
+		return nil, "", ErrNotFound
+	}
+	return innerNode, lastPart, nil
 }
 
 func (c *ntmConfig) addToKnownKeys(key string) {
@@ -609,7 +681,7 @@ func (c *ntmConfig) AllSettingsWithoutDefault() map[string]interface{} {
 	defer c.RUnlock()
 
 	// We only want to include leaf with a source higher than SourceDefault
-	return c.root.DumpSettings(func(source model.Source) bool { return source.IsGreaterOrEqualThan(model.SourceUnknown) })
+	return c.root.DumpSettings(func(source model.Source) bool { return source.IsGreaterThan(model.SourceDefault) })
 }
 
 // AllSettingsBySource returns the settings from each source (file, env vars, ...)

pkg/config/nodetreemodel/config_test.go

@@ -433,3 +433,145 @@ server_timeout
   val:30, source:default`
 	assert.Equal(t, expect, txt)
 }
+
+func TestUnsetForSource(t *testing.T) {
+	// env source, highest priority
+	os.Setenv("TEST_NETWORK_PATH_COLLECTOR_INPUT_CHAN_SIZE", "23456")
+	os.Setenv("TEST_NETWORK_PATH_COLLECTOR_PATHTEST_CONTEXTS_LIMIT", "654321")
+	os.Setenv("TEST_NETWORK_PATH_COLLECTOR_PROCESSING_CHAN_SIZE", "78900")
+	// file source, medium priority
+	configData := `network_path:
+  collector:
+    workers: 6
+    pathtest_contexts_limit: 43210
+    processing_chan_size: 45678`
+	// default source, lowest priority
+	cfg := NewConfig("test", "TEST", strings.NewReplacer(".", "_"))
+	cfg.BindEnvAndSetDefault("network_path.collector.input_chan_size", 100000)
+	cfg.BindEnvAndSetDefault("network_path.collector.pathtest_contexts_limit", 100000)
+	cfg.BindEnvAndSetDefault("network_path.collector.processing_chan_size", 100000)
+	cfg.BindEnvAndSetDefault("network_path.collector.workers", 4)
+
+	cfg.BuildSchema()
+	err := cfg.ReadConfig(strings.NewReader(configData))
+	require.NoError(t, err)
+
+	// The merged config
+	txt := cfg.(*ntmConfig).Stringify("root")
+	expect := `network_path
+  collector
+    input_chan_size
+      val:23456, source:environment-variable
+    pathtest_contexts_limit
+      val:654321, source:environment-variable
+    processing_chan_size
+      val:78900, source:environment-variable
+    workers
+      val:6, source:file`
+	assert.Equal(t, expect, txt)
+
+	// No change if source doesn't match
+	cfg.UnsetForSource("network_path.collector.input_chan_size", model.SourceFile)
+	assert.Equal(t, expect, txt)
+
+	// No change if setting is not a leaf
+	cfg.UnsetForSource("network_path", model.SourceEnvVar)
+	assert.Equal(t, expect, txt)
+
+	// No change if setting is not found
+	cfg.UnsetForSource("network_path.unknown", model.SourceEnvVar)
+	assert.Equal(t, expect, txt)
+
+	// Remove a setting from the env source, nothing in the file source, it goes to default
+	cfg.UnsetForSource("network_path.collector.input_chan_size", model.SourceEnvVar)
+	txt = cfg.(*ntmConfig).Stringify("root")
+	expect = `network_path
+  collector
+    input_chan_size
+      val:100000, source:default
+    pathtest_contexts_limit
+      val:654321, source:environment-variable
+    processing_chan_size
+      val:78900, source:environment-variable
+    workers
+      val:6, source:file`
+	assert.Equal(t, expect, txt)
+
+	// Remove a setting from the file source, it goes to default
+	cfg.UnsetForSource("network_path.collector.workers", model.SourceFile)
+	txt = cfg.(*ntmConfig).Stringify("root")
+	expect = `network_path
+  collector
+    input_chan_size
+      val:100000, source:default
+    pathtest_contexts_limit
+      val:654321, source:environment-variable
+    processing_chan_size
+      val:78900, source:environment-variable
+    workers
+      val:4, source:default`
+	assert.Equal(t, expect, txt)
+
+	// Removing a setting from the env source, it goes to file source
+	cfg.UnsetForSource("network_path.collector.processing_chan_size", model.SourceEnvVar)
+	txt = cfg.(*ntmConfig).Stringify("root")
+	expect = `network_path
+  collector
+    input_chan_size
+      val:100000, source:default
+    pathtest_contexts_limit
+      val:654321, source:environment-variable
+    processing_chan_size
+      val:45678, source:file
+    workers
+      val:4, source:default`
+	assert.Equal(t, expect, txt)
+
+	// Then remove it from the file source as well, leaving the default source
+	cfg.UnsetForSource("network_path.collector.processing_chan_size", model.SourceFile)
+	txt = cfg.(*ntmConfig).Stringify("root")
+	expect = `network_path
+  collector
+    input_chan_size
+      val:100000, source:default
+    pathtest_contexts_limit
+      val:654321, source:environment-variable
+    processing_chan_size
+      val:100000, source:default
+    workers
+      val:4, source:default`
+	assert.Equal(t, expect, txt)
+
+	// Check the file layer in isolation
+	fileTxt := cfg.(*ntmConfig).Stringify(model.SourceFile)
+	fileExpect := `network_path
+  collector
+    pathtest_contexts_limit
+      val:43210, source:file`
+	assert.Equal(t, fileExpect, fileTxt)
+
+	// Removing from the file source first does not change the merged value, because it uses env layer
+	cfg.UnsetForSource("network_path.collector.pathtest_contexts_limit", model.SourceFile)
+	assert.Equal(t, expect, txt)
+
+	// But the file layer itself has been modified
+	fileTxt = cfg.(*ntmConfig).Stringify(model.SourceFile)
+	fileExpect = `network_path
+  collector`
+	assert.Equal(t, fileExpect, fileTxt)
+
+	// Finally, remove it from the env layer
+	cfg.UnsetForSource("network_path.collector.pathtest_contexts_limit", model.SourceEnvVar)
+	txt = cfg.(*ntmConfig).Stringify("root")
+	expect = `network_path
+  collector
+    input_chan_size
+      val:100000, source:default
+    pathtest_contexts_limit
+      val:100000, source:default
+    processing_chan_size
+      val:100000, source:default
+    workers
+      val:4, source:default`
+	assert.Equal(t, expect, txt)
+}

pkg/config/nodetreemodel/inner_node.go

@@ -88,7 +88,7 @@ func (n *innerNode) Merge(src InnerNode) error {
 			}
 
 			if srcIsLeaf {
-				if srcLeaf.SourceGreaterOrEqual(dstLeaf.Source()) {
+				if srcLeaf.Source() == dstLeaf.Source() || srcLeaf.SourceGreaterThan(dstLeaf.Source()) {
 					n.children[name] = srcLeaf.Clone()
 				}
 			} else {
@@ -132,7 +132,7 @@ func (n *innerNode) SetAt(key []string, value interface{}, source model.Source)
 		}
 
 		if leaf, ok := node.(LeafNode); ok {
-			if source.IsGreaterOrEqualThan(leaf.Source()) {
+			if source == leaf.Source() || source.IsGreaterThan(leaf.Source()) {
 				n.children[part] = newLeafNode(value, source)
 				return true, nil
 			}
@@ -163,6 +163,12 @@ func (n *innerNode) InsertChildNode(name string, node Node) {
 	n.makeRemapCase()
 }
 
+// RemoveChild removes a node from the current node
+func (n *innerNode) RemoveChild(name string) {
+	delete(n.children, name)
+	n.makeRemapCase()
+}
+
 // DumpSettings clone the entire tree starting from the node into a map based on the leaf source.
 //
 // The selector will be call with the source of each leaf to determine if it should be included in the dump.

pkg/config/nodetreemodel/leaf_node.go

@@ -43,10 +43,10 @@ func (n *leafNodeImpl) Clone() Node {
 	return newLeafNode(n.val, n.source)
 }
 
-// SourceGreaterOrEqual returns true if the source of the current node is greater or equal to the one given as a
+// SourceGreaterThan returns true if the source of the current node is greater than the one given as a
 // parameter
-func (n *leafNodeImpl) SourceGreaterOrEqual(source model.Source) bool {
-	return n.source.IsGreaterOrEqualThan(source)
+func (n *leafNodeImpl) SourceGreaterThan(source model.Source) bool {
+	return n.source.IsGreaterThan(source)
 }
 
 // GetChild returns an error because a leaf has no children

pkg/config/nodetreemodel/missing_node.go

@@ -38,6 +38,6 @@ func (m *missingLeafImpl) Clone() Node {
 	return m
 }
 
-func (m *missingLeafImpl) SourceGreaterOrEqual(model.Source) bool {
+func (m *missingLeafImpl) SourceGreaterThan(model.Source) bool {
 	return false
 }

pkg/config/nodetreemodel/node.go

@@ -96,6 +96,7 @@ type InnerNode interface {
 	Merge(InnerNode) error
 	SetAt([]string, interface{}, model.Source) (bool, error)
 	InsertChildNode(string, Node)
+	RemoveChild(string)
 	makeRemapCase()
 	DumpSettings(func(model.Source) bool) map[string]interface{}
 }
@@ -105,5 +106,5 @@ type LeafNode interface {
 	Node
 	Get() interface{}
 	Source() model.Source
-	SourceGreaterOrEqual(model.Source) bool
+	SourceGreaterThan(model.Source) bool
 }

pkg/config/nodetreemodel/struct_node.go

@@ -74,6 +74,9 @@ func (n *structNodeImpl) SetAt([]string, interface{}, model.Source) (bool, error
 // InsertChildNode is not implemented for a leaf node
 func (n *structNodeImpl) InsertChildNode(string, Node) {}
 
+// RemoveChild is not implemented for struct node
+func (n *structNodeImpl) RemoveChild(string) {}
+
 // makeRemapCase not implemented
 func (n *structNodeImpl) makeRemapCase() {}