Drmorr/custom flush times

.github/FUNDING.yml

@@ -1 +1,2 @@
+---
 github: acrlabs

.github/workflows/verify.yml

@@ -1,5 +1,6 @@
+---
 name: code verification
-on: [push]
+on: [push]  # yamllint disable-line rule:truthy
 
 jobs:
   build:
@@ -27,10 +28,8 @@ jobs:
         uses: actions/setup-go@v4
         with:
           go-version-file: go.mod
-      - name: Lint Go code
-        uses: golangci/golangci-lint-action@v3
-        with:
-          version: v1.55
+      - name: Run pre-commit
+        uses: pre-commit/[email protected]
       - name: Test
         run: make test cover
       - name: Upload coverage

.golangci.yml

@@ -1,3 +1,4 @@
+---
 run:
   timeout: 5m
 
@@ -53,7 +54,7 @@ linters-settings:
 
   gosec:
     excludes:
-      - G601 # memory aliasing -- not a problem in 1.22+
+      - G601  # memory aliasing -- not a problem in 1.22+
 
   importas:
     no-unaliased: true

.pre-commit-config.yaml

@@ -1,13 +1,16 @@
 ---
 repos:
-  - repo: https://github.com/pre-commit/pre-commit-hooks
+  - repo: "https://github.com/pre-commit/pre-commit-hooks"
     rev: v4.4.0
     hooks:
       - id: end-of-file-fixer
-      - id: check-yaml
-        args: ["--allow-multiple-documents"]
       - id: trailing-whitespace
-  - repo: https://github.com/golangci/golangci-lint
+  - repo: "https://github.com/adrienverge/yamllint"
+    rev: v1.35.1
+    hooks:
+      - id: yamllint
+        args: ['--strict']
+  - repo: "https://github.com/golangci/golangci-lint"
     rev: v1.55.2
     hooks:
       - id: golangci-lint

.yamllint

@@ -0,0 +1,9 @@
+---
+extends: default
+
+rules:
+  line-length:
+    max: 120
+  quoted-strings:
+    required: false
+    extra-required: [^http://, ^https://]

Makefile

@@ -5,8 +5,8 @@ GO_COVER_FILE=$(COVERAGE_DIR)/go-coverage.txt
 include build/base.mk
 include build/k8s.mk
 
-$(ARTIFACTS)::
-	CGO_ENABLED=0 go build -trimpath -o $(BUILD_DIR)/$@ ./cmd/.
+main:
+	CGO_ENABLED=0 go build -ldflags "-s -w" -trimpath -o $(BUILD_DIR)/prom2parquet ./cmd/.
 
 lint:
 	golangci-lint run

cmd/args.go

@@ -3,18 +3,20 @@ package main
 import (
 	"sort"
 	"strings"
+	"time"
 
 	log "github.com/sirupsen/logrus"
 
 	"github.com/acrlabs/prom2parquet/pkg/backends"
 )
 
 const (
-	prefixFlag      = "prefix"
-	serverPortFlag  = "server-port"
-	backendFlag     = "backend"
-	backendRootFlag = "backend-root"
-	verbosityFlag   = "verbosity"
+	prefixFlag        = "prefix"
+	serverPortFlag    = "server-port"
+	flushIntervalFlag = "flush-interval"
+	backendFlag       = "backend"
+	backendRootFlag   = "backend-root"
+	verbosityFlag     = "verbosity"
 )
 
 //nolint:gochecknoglobals
@@ -35,10 +37,10 @@ var logLevelIDs = map[log.Level][]string{
 }
 
 type options struct {
-	port int
-
-	backend     backends.StorageBackend
-	backendRoot string
+	port          int
+	flushInterval time.Duration
+	backend       backends.StorageBackend
+	backendRoot   string
 
 	verbosity log.Level
 }

cmd/root.go

@@ -3,6 +3,7 @@
 import (
 	"fmt"
 	"os"
+	"time"
 
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
@@ -26,6 +27,13 @@
 		},
 	}
 
+	root.PersistentFlags().DurationVar(
+		&opts.flushInterval,
+		flushIntervalFlag,
+		10*time.Minute,
+		"data flush interval",
+	)
+
 	root.PersistentFlags().IntVarP(
 		&opts.port,
 		serverPortFlag,

cmd/server.go

@@ -149,6 +149,7 @@ func (self *promserver) spawnWriter(ctx context.Context, channelName string) (ch
 		self.opts.backendRoot,
 		channelName,
 		self.opts.backend,
+		self.opts.flushInterval,
 	)
 	if err != nil {
 		return nil, fmt.Errorf("could not create writer for %s: %w", channelName, err)

k8s/main.py

@@ -27,6 +27,8 @@ def __init__(self):
             image=image,
             args=[
                 "/prom2parquet",
+                "--backend", "s3",
+                "--backend-root", "simkube",
             ],
         ).with_env(env).with_ports(SERVER_PORT).with_security_context(Capability.DEBUG)
 

pkg/parquet/writer.go

@@ -17,13 +17,13 @@
 const pageNum = 4
 
 type Prom2ParquetWriter struct {
-	backend backends.StorageBackend
-	root    string
-	prefix  string
+	backend       backends.StorageBackend
+	root          string
+	prefix        string
+	flushInterval time.Duration
 
-	currentFile   string
-	nextFlushTime time.Time
-	pw            *writer.ParquetWriter
+	currentFile string
+	pw          *writer.ParquetWriter
 
 	clock clockwork.Clock
 }
@@ -43,30 +43,50 @@
 	ctx context.Context,
 	root, prefix string,
 	backend backends.StorageBackend,
+	flushInterval time.Duration,
 ) (*Prom2ParquetWriter, error) {
 	return &Prom2ParquetWriter{
-		backend: backend,
-		root:    root,
-		prefix:  prefix,
+		backend:       backend,
+		root:          root,
+		prefix:        prefix,
+		flushInterval: flushInterval,
 
 		clock: clockwork.NewRealClock(),
 	}, nil
 }
 
 func (self *Prom2ParquetWriter) Listen(stream <-chan prompb.TimeSeries) {
-	if err := self.flush(); err != nil {
-		log.Errorf("could not flush writer: %v", err)
+	self.listen(stream, self.getFlushTimer(), nil)
+}
+
+func (self *Prom2ParquetWriter) listen(
+	stream <-chan prompb.TimeSeries,
+	flushTimer <-chan time.Time,
+	running chan<- bool, // used for testing
+) {
+	if err := self.createBackendWriter(); err != nil {
+		log.Errorf("could not create backend writer: %v", err)
 		return
 	}
-	defer self.closeFile()
 
-	flushTicker := time.NewTicker(time.Minute)
+	// self.pw is a pointer to the writer instance, but it can get switched
+	// out from under us whenever we flush; go defer evaluates the function
+	// args when the defer call happens, not when the deferred function actually
+	// executes, so here we need to use a double pointer so that we can make
+	// sure we're closing the actual correct writer instance
+	defer func(pw **writer.ParquetWriter) {
+		closeFile(*pw)
+		close(running)
+	}(&self.pw)
+
+	if running != nil {
+		running <- true
+	}
 
 	for {
 		select {
 		case ts, ok := <-stream:
 			if !ok {
-				self.closeFile()
 				return
 			}
 
@@ -79,36 +99,29 @@
 					log.Errorf("could not write datapoint: %v", err)
 				}
 			}
-		case <-flushTicker.C:
-			if time.Now().After(self.nextFlushTime) {
-				log.Infof("Flush triggered: %v >= %v", time.Now(), self.nextFlushTime)
-				if err := self.flush(); err != nil {
-					log.Errorf("could not flush data: %v", err)
-					return
-				}
+		case <-flushTimer:
+			flushTimer = self.getFlushTimer()
+			log.Infof("flush triggered for %v", self.currentFile)
+
+			// Run this in a separate goroutine so that writing the data
+			// to S3 (with throttling or whatever) doesn't block the new incoming
+			// datapoints
+			go closeFile(self.pw)
+			if err := self.createBackendWriter(); err != nil {
+				log.Errorf("could not create backend writer: %v", err)
+				return
 			}
 		}
 	}
 }
 
-func (self *Prom2ParquetWriter) closeFile() {
-	if self.pw != nil {
-		if err := self.pw.WriteStop(); err != nil {
-			log.Errorf("can't close parquet writer: %v", err)
-		}
-		self.pw = nil
-	}
-}
-
-func (self *Prom2ParquetWriter) flush() error {
-	now := self.clock.Now().UTC()
-
-	self.closeFile()
+func (self *Prom2ParquetWriter) createBackendWriter() error {
+	basename := self.now().Truncate(self.flushInterval).Format("20060102150405")
+	self.currentFile = fmt.Sprintf("%s/%s.parquet", self.prefix, basename)
 
-	self.currentFile = fmt.Sprintf("%s/%s.parquet", self.prefix, now.Format("2006010215"))
 	fw, err := backends.ConstructBackendForFile(self.root, self.currentFile, self.backend)
 	if err != nil {
-		return fmt.Errorf("can't create storage backend writer: %w", err)
+		return fmt.Errorf("can't create storage backend: %w", err)
 	}
 
 	pw, err := writer.NewParquetWriter(fw, new(DataPoint), pageNum)
@@ -118,21 +131,24 @@
 	pw.CompressionType = parquet.CompressionCodec_SNAPPY
 
 	self.pw = pw
-	self.advanceFlushTime(&now)
 
 	return nil
 }
 
-func (self *Prom2ParquetWriter) advanceFlushTime(now *time.Time) {
-	nextHour := now.Add(time.Hour)
-	self.nextFlushTime = time.Date(
-		nextHour.Year(),
-		nextHour.Month(),
-		nextHour.Day(),
-		nextHour.Hour(),
-		0,
-		0,
-		0,
-		nextHour.Location(),
-	)
+func (self *Prom2ParquetWriter) getFlushTimer() <-chan time.Time {
+	now := self.now()
+	nextFlushTime := now.Truncate(self.flushInterval).Add(self.flushInterval)
+	return time.After(nextFlushTime.Sub(now))
+}
+
+func (self *Prom2ParquetWriter) now() time.Time {
+	return self.clock.Now().UTC()
+}
+
+func closeFile(pw *writer.ParquetWriter) {
+	if pw != nil {
+		if err := pw.WriteStop(); err != nil {
+			log.Errorf("can't close parquet writer: %v", err)
+		}
+	}
 }