Convert to using prometheus appendable

component/prometheus/fanout.go

@@ -0,0 +1,132 @@
+package prometheus
+
+import (
+	"context"
+	"sync"
+
+	"github.com/prometheus/prometheus/model/exemplar"
+	"github.com/prometheus/prometheus/model/labels"
+	"github.com/prometheus/prometheus/model/metadata"
+
+	"github.com/prometheus/prometheus/storage"
+)
+
+var _ storage.Appendable = (*Fanout)(nil)
+
+type intercept func(ref storage.SeriesRef, l labels.Labels, t int64, v float64) (storage.SeriesRef, labels.Labels, int64, float64, error)
+
+// Fanout supports the default Flow style of appendables since it can go to multiple outputs. It also allows the intercepting of appends.
+type Fanout struct {
+	mut sync.RWMutex
+	// intercept allows one to intercept the series before it fans out to make any changes. If labels.Labels returns nil the series is not propagated.
+	// Intercept shouuld be thread safe and can be called across appenders.
+	intercept func(ref storage.SeriesRef, l labels.Labels, t int64, v float64) (storage.SeriesRef, labels.Labels, int64, float64, error)
+
+	// children is where to fan out.
+	children []storage.Appendable
+
+	// ComponentID is what component this belongs to.
+	componentID string
+}
+
+// NewFanout creates a fanout appendable.
+func NewFanout(inter intercept, children []storage.Appendable, componentID string) *Fanout {
+	return &Fanout{
+		intercept:   inter,
+		children:    children,
+		componentID: componentID,
+	}
+}
+
+// UpdateChildren allows changing of the children of the fanout.
+func (f *Fanout) UpdateChildren(children []storage.Appendable) {
+	f.mut.Lock()
+	defer f.mut.Unlock()
+	f.children = children
+}
+
+// Appender satisfies the Appendable interface.
+func (f *Fanout) Appender(ctx context.Context) storage.Appender {
+	f.mut.RLock()
+	defer f.mut.RUnlock()
+
+	app := &appender{
+		children:    make([]storage.Appender, len(f.children)),
+		intercept:   f.intercept,
+		componentID: f.componentID,
+	}
+	for i, x := range f.children {
+		if x == nil {
+			continue
+		}
+		app.children[i] = x.Appender(ctx)
+	}
+	return app
+}
+
+var _ storage.Appender = (*appender)(nil)
+
+type appender struct {
+	children    []storage.Appender
+	componentID string
+	intercept   func(ref storage.SeriesRef, l labels.Labels, t int64, v float64) (storage.SeriesRef, labels.Labels, int64, float64, error)
+}
+
+// Append satisfies the Appender interface.
+func (a *appender) Append(ref storage.SeriesRef, l labels.Labels, t int64, v float64) (storage.SeriesRef, error) {
+	if ref == 0 {
+		ref = storage.SeriesRef(GlobalRefMapping.GetOrAddGlobalRefID(l))
+	}
+	nRef := ref
+	nL := l
+	nT := t
+	nV := v
+	if a.intercept != nil {
+		var err error
+		nRef, nL, nT, nV, err = a.intercept(ref, l, t, v)
+		if err != nil {
+			return 0, err
+		}
+	}
+	for _, x := range a.children {
+		if x == nil || nL == nil {
+			continue
+		}
+		_, _ = x.Append(nRef, nL, nT, nV)
+	}
+	return ref, nil
+}
+
+// Commit satisfies the Appender interface.
+func (a *appender) Commit() error {
+	for _, x := range a.children {
+		if x == nil {
+			continue
+		}
+		_ = x.Commit()
+	}
+	return nil
+}
+
+// Rollback satisifies the Appender interface.
+func (a *appender) Rollback() error {
+	for _, x := range a.children {
+		if x == nil {
+			continue
+		}
+		_, _ = x, a.Rollback()
+	}
+	return nil
+}
+
+// AppendExemplar satisfies the Appender interface.
+func (a *appender) AppendExemplar(ref storage.SeriesRef, l labels.Labels, e exemplar.Exemplar) (storage.SeriesRef, error) {
+	// TODO implement me
+	panic("implement me")
+}
+
+// UpdateMetadata satisifies the Appender interface.
+func (a *appender) UpdateMetadata(ref storage.SeriesRef, l labels.Labels, m metadata.Metadata) (storage.SeriesRef, error) {
+	// TODO implement me
+	panic("implement me")
+}

component/prometheus/globalrefmap.go

@@ -44,7 +44,7 @@ func newGlobalRefMap() *GlobalRefMap {
 }
 
 // GetOrAddLink is called by a remote_write endpoint component to add mapping and get back the global id.
-func (g *GlobalRefMap) GetOrAddLink(componentID string, localRefID uint64, fm *FlowMetric) uint64 {
+func (g *GlobalRefMap) GetOrAddLink(componentID string, localRefID uint64, lbls labels.Labels) uint64 {
 	g.mut.Lock()
 	defer g.mut.Unlock()
 
@@ -59,7 +59,7 @@ func (g *GlobalRefMap) GetOrAddLink(componentID string, localRefID uint64, fm *F
 		g.mappings[componentID] = m
 	}
 
-	labelHash := fm.labels.Hash()
+	labelHash := lbls.Hash()
 	globalID, found := g.labelsHashToGlobal[labelHash]
 	if found {
 		m.localToGlobal[localRefID] = globalID
@@ -79,6 +79,10 @@ func (g *GlobalRefMap) GetOrAddGlobalRefID(l labels.Labels) uint64 {
 	g.mut.Lock()
 	defer g.mut.Unlock()
 
+	if l == nil {
+		return 0
+	}
+
 	labelHash := l.Hash()
 	globalID, found := g.labelsHashToGlobal[labelHash]
 	if found {

component/prometheus/globalrefmap_test.go

@@ -48,8 +48,7 @@ func TestAddingLocalMapping(t *testing.T) {
 	})
 
 	globalID := mapping.GetOrAddGlobalRefID(l)
-	fm := NewFlowMetric(globalID, l, 0)
-	shouldBeSameGlobalID := mapping.GetOrAddLink("1", 1, fm)
+	shouldBeSameGlobalID := mapping.GetOrAddLink("1", 1, l)
 	require.True(t, globalID == shouldBeSameGlobalID)
 	require.Len(t, mapping.labelsHashToGlobal, 1)
 	require.Len(t, mapping.mappings, 1)
@@ -67,9 +66,8 @@ func TestAddingLocalMappings(t *testing.T) {
 	})
 
 	globalID := mapping.GetOrAddGlobalRefID(l)
-	fm := NewFlowMetric(globalID, l, 0)
-	shouldBeSameGlobalID := mapping.GetOrAddLink("1", 1, fm)
-	shouldBeSameGlobalID2 := mapping.GetOrAddLink("2", 1, fm)
+	shouldBeSameGlobalID := mapping.GetOrAddLink("1", 1, l)
+	shouldBeSameGlobalID2 := mapping.GetOrAddLink("2", 1, l)
 	require.True(t, globalID == shouldBeSameGlobalID)
 	require.True(t, globalID == shouldBeSameGlobalID2)
 	require.Len(t, mapping.labelsHashToGlobal, 1)
@@ -92,10 +90,8 @@ func TestAddingLocalMappingsWithoutCreatingGlobalUpfront(t *testing.T) {
 		Value: "test",
 	})
 
-	fm := NewFlowMetric(1, l, 0)
-
-	shouldBeSameGlobalID := mapping.GetOrAddLink("1", 1, fm)
-	shouldBeSameGlobalID2 := mapping.GetOrAddLink("2", 1, fm)
+	shouldBeSameGlobalID := mapping.GetOrAddLink("1", 1, l)
+	shouldBeSameGlobalID2 := mapping.GetOrAddLink("2", 1, l)
 	require.True(t, shouldBeSameGlobalID2 == shouldBeSameGlobalID)
 	require.Len(t, mapping.labelsHashToGlobal, 1)
 	require.Len(t, mapping.mappings, 2)
@@ -122,11 +118,8 @@ func TestStaleness(t *testing.T) {
 		Value: "test2",
 	})
 
-	fm := NewFlowMetric(0, l, 0)
-	fm2 := NewFlowMetric(0, l2, 0)
-
-	global1 := mapping.GetOrAddLink("1", 1, fm)
-	_ = mapping.GetOrAddLink("2", 1, fm2)
+	global1 := mapping.GetOrAddLink("1", 1, l)
+	_ = mapping.GetOrAddLink("2", 1, l2)
 	mapping.AddStaleMarker(global1, l)
 	require.Len(t, mapping.staleGlobals, 1)
 	require.Len(t, mapping.labelsHashToGlobal, 2)
@@ -145,9 +138,7 @@ func TestRemovingStaleness(t *testing.T) {
 		Value: "test",
 	})
 
-	fm := NewFlowMetric(0, l, 0)
-
-	global1 := mapping.GetOrAddLink("1", 1, fm)
+	global1 := mapping.GetOrAddLink("1", 1, l)
 	mapping.AddStaleMarker(global1, l)
 	require.Len(t, mapping.staleGlobals, 1)
 	mapping.RemoveStaleMarker(global1)