CostModel.ComputeAllocation: rewrite queryPods to work in batches and process results into a pod-based mapping

pkg/kubecost/window.go

@@ -8,6 +8,8 @@ import (
 	"strconv"
 	"time"
 
+	"github.com/kubecost/cost-model/pkg/env"
+	"github.com/kubecost/cost-model/pkg/thanos"
 	"github.com/kubecost/cost-model/pkg/util"
 )
 
@@ -574,6 +576,47 @@ func (w Window) DurationOffset() (time.Duration, time.Duration, error) {
 	return duration, offset, nil
 }
 
+// DurationOffsetForPrometheus returns durations representing the duration and
+// offset of the given window, factoring in the Thanos offset if necessary. The
+// duration is returned as
+func (w Window) DurationOffsetForPrometheus() (string, string, error) {
+	duration, offset, err := w.DurationOffset()
+	if err != nil {
+		return "", "", err
+	}
+
+	// If using Thanos, increase offset to 3 hours, reducing the duration by
+	// equal measure to maintain the same starting point.
+	thanosDur := thanos.OffsetDuration()
+	if offset < thanosDur && env.IsThanosEnabled() {
+		diff := thanosDur - offset
+		offset += diff
+		duration -= diff
+	}
+
+	// If duration < 0, return an error
+	if duration < 0 {
+		return "", "", fmt.Errorf("negative duration: %s", duration)
+	}
+
+	// Negative offset means that the end time is in the future. Prometheus
+	// fails for non-positive offset values, so shrink the duration and
+	// remove the offset altogether.
+	if offset < 0 {
+		duration = duration + offset
+		offset = 0
+	}
+
+	durStr, offStr := util.DurationOffsetStrings(duration, offset)
+	if offset < time.Minute {
+		offStr = ""
+	} else {
+		offStr = " offset " + offStr
+	}
+
+	return durStr, offStr, nil
+}
+
 // DurationOffsetStrings returns formatted, Prometheus-compatible strings representing
 // the duration and offset of the window in terms of days, hours, minutes, or seconds;
 // e.g. ("7d", "1441m", "30m", "1s", "")

pkg/kubecost/window_test.go

@@ -2,8 +2,11 @@ package kubecost
 
 import (
 	"fmt"
+	"strings"
 	"testing"
 	"time"
+
+	"github.com/kubecost/cost-model/pkg/env"
 )
 
 func TestRoundBack(t *testing.T) {
@@ -601,6 +604,162 @@ func TestWindow_DurationOffsetStrings(t *testing.T) {
 	}
 }
 
+func TestWindow_DurationOffsetForPrometheus(t *testing.T) {
+	// Set-up and tear-down
+	thanosEnabled := env.GetBool(env.ThanosEnabledEnvVar, false)
+	defer env.SetBool(env.ThanosEnabledEnvVar, thanosEnabled)
+
+	// Test for Prometheus (env.IsThanosEnabled() == false)
+	env.SetBool(env.ThanosEnabledEnvVar, false)
+	if env.IsThanosEnabled() {
+		t.Fatalf("expected env.IsThanosEnabled() == false")
+	}
+
+	w, err := ParseWindowUTC("1d")
+	if err != nil {
+		t.Fatalf(`unexpected error parsing "1d": %s`, err)
+	}
+	dur, off, err := w.DurationOffsetForPrometheus()
+	if err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+	if dur != "1d" {
+		t.Fatalf(`expect: window to be "1d"; actual: "%s"`, dur)
+	}
+	if off != "" {
+		t.Fatalf(`expect: offset to be ""; actual: "%s"`, off)
+	}
+
+	w, err = ParseWindowUTC("2h")
+	if err != nil {
+		t.Fatalf(`unexpected error parsing "2h": %s`, err)
+	}
+	dur, off, err = w.DurationOffsetForPrometheus()
+	if err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+	if dur != "2h" {
+		t.Fatalf(`expect: window to be "2h"; actual: "%s"`, dur)
+	}
+	if off != "" {
+		t.Fatalf(`expect: offset to be ""; actual: "%s"`, off)
+	}
+
+	w, err = ParseWindowUTC("10m")
+	if err != nil {
+		t.Fatalf(`unexpected error parsing "10m": %s`, err)
+	}
+	dur, off, err = w.DurationOffsetForPrometheus()
+	if err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+	if dur != "10m" {
+		t.Fatalf(`expect: window to be "10m"; actual: "%s"`, dur)
+	}
+	if off != "" {
+		t.Fatalf(`expect: offset to be ""; actual: "%s"`, off)
+	}
+
+	w, err = ParseWindowUTC("1589448338,1589534798")
+	if err != nil {
+		t.Fatalf(`unexpected error parsing "1589448338,1589534798": %s`, err)
+	}
+	dur, off, err = w.DurationOffsetForPrometheus()
+	if err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+	if dur != "1441m" {
+		t.Fatalf(`expect: window to be "1441m"; actual: "%s"`, dur)
+	}
+	if !strings.HasPrefix(off, " offset ") {
+		t.Fatalf(`expect: offset to start with " offset "; actual: "%s"`, off)
+	}
+
+	w, err = ParseWindowUTC("yesterday")
+	if err != nil {
+		t.Fatalf(`unexpected error parsing "yesterday": %s`, err)
+	}
+	dur, off, err = w.DurationOffsetForPrometheus()
+	if err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+	if dur != "1d" {
+		t.Fatalf(`expect: window to be "1d"; actual: "%s"`, dur)
+	}
+	if !strings.HasPrefix(off, " offset ") {
+		t.Fatalf(`expect: offset to start with " offset "; actual: "%s"`, off)
+	}
+
+	// Test for Thanos (env.IsThanosEnabled() == true)
+	env.SetBool(env.ThanosEnabledEnvVar, true)
+	if !env.IsThanosEnabled() {
+		t.Fatalf("expected env.IsThanosEnabled() == true")
+	}
+
+	w, err = ParseWindowUTC("1d")
+	if err != nil {
+		t.Fatalf(`unexpected error parsing "1d": %s`, err)
+	}
+	dur, off, err = w.DurationOffsetForPrometheus()
+	if err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+	if dur != "21h" {
+		t.Fatalf(`expect: window to be "21d"; actual: "%s"`, dur)
+	}
+	if off != " offset 3h" {
+		t.Fatalf(`expect: offset to be " offset 3h"; actual: "%s"`, off)
+	}
+
+	w, err = ParseWindowUTC("2h")
+	if err != nil {
+		t.Fatalf(`unexpected error parsing "2h": %s`, err)
+	}
+	dur, off, err = w.DurationOffsetForPrometheus()
+	if err == nil {
+		t.Fatalf(`expected error (negative duration); got ("%s", "%s")`, dur, off)
+	}
+
+	w, err = ParseWindowUTC("10m")
+	if err != nil {
+		t.Fatalf(`unexpected error parsing "1d": %s`, err)
+	}
+	dur, off, err = w.DurationOffsetForPrometheus()
+	if err == nil {
+		t.Fatalf(`expected error (negative duration); got ("%s", "%s")`, dur, off)
+	}
+
+	w, err = ParseWindowUTC("1589448338,1589534798")
+	if err != nil {
+		t.Fatalf(`unexpected error parsing "1589448338,1589534798": %s`, err)
+	}
+	dur, off, err = w.DurationOffsetForPrometheus()
+	if err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+	if dur != "1441m" {
+		t.Fatalf(`expect: window to be "1441m"; actual: "%s"`, dur)
+	}
+	if !strings.HasPrefix(off, " offset ") {
+		t.Fatalf(`expect: offset to start with " offset "; actual: "%s"`, off)
+	}
+
+	w, err = ParseWindowUTC("yesterday")
+	if err != nil {
+		t.Fatalf(`unexpected error parsing "yesterday": %s`, err)
+	}
+	dur, off, err = w.DurationOffsetForPrometheus()
+	if err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+	if dur != "1d" {
+		t.Fatalf(`expect: window to be "1d"; actual: "%s"`, dur)
+	}
+	if !strings.HasPrefix(off, " offset ") {
+		t.Fatalf(`expect: offset to start with " offset "; actual: "%s"`, off)
+	}
+}
+
 // TODO
 // func TestWindow_Overlaps(t *testing.T) {}