lib/protoparser/native: follow-up after fe01f4803d

docs/CHANGELOG.md

@@ -17,6 +17,7 @@ The following tip changes can be tested by building VictoriaMetrics components f
 
 FEATURE: [vmalert](https://docs.victoriametrics.com/vmalert.html): add support for `alert_relabel_configs` option at `-notifier.config`. This option allows configuring relabeling rules for alerts before sending them to configured notifiers. See [these docs](https://docs.victoriametrics.com/vmalert.html#notifier-configuration-file) for details.
 
+BUGFIX: fix goroutine leak and possible deadlock when importing invalid data via [native binary format](https://docs.victoriametrics.com/#how-to-import-data-in-native-format). See [this pull request](https://github.com/VictoriaMetrics/VictoriaMetrics/pull/2423).
 BUGFIX: [Graphite Render API](https://docs.victoriametrics.com/#graphite-render-api-usage): properly calculate [hitCount](https://graphite.readthedocs.io/en/latest/functions.html#graphite.render.functions.hitcount) function. Previously it could return empty results if there were no original samples in some parts of the selected time range.
 BUGFIX: [MetricsQL](https://docs.victoriametrics.com/MetricsQL.html): allow overriding built-in function names inside [WITH templates](https://play.victoriametrics.com/promql/expand-with-exprs). For example, `WITH (sum(a,b) = a + b + 1) sum(x,y)` now expands into `x + y + 1`. Previously such a query would fail with `cannot use reserved name` error. See [this bugreport](https://github.com/VictoriaMetrics/metricsql/issues/5).
 

lib/protoparser/native/streamparser.go

@@ -42,57 +42,37 @@ func ParseStream(r io.Reader, isGzip bool, callback func(block *Block) error) er
 
 	// Read native blocks and feed workers with work.
 	sizeBuf := make([]byte, 4)
-	var wg sync.WaitGroup
+
-	var (
+	ctx := &streamContext{}
-		callbackErrLock sync.Mutex
-		callbackErr     error
-	)
-	addFirstErr := func(err error) {
-		processErrors.Inc()
-		callbackErrLock.Lock()
-		if callbackErr == nil {
-			callbackErr = fmt.Errorf("error when processing native block: %w", err)
-		}
-		callbackErrLock.Unlock()
-	}
 	for {
 		uw := getUnmarshalWork()
 		uw.tr = tr
-		uw.callback = func(block *Block, parseErr error) {
+		uw.ctx = ctx
-			defer wg.Done()
+		uw.callback = callback
-			// fast path
-			if parseErr != nil {
-				addFirstErr(parseErr)
-				return
-			}
-			if err := callback(block); err != nil {
-				addFirstErr(err)
-			}
-		}
 
 		// Read uw.metricNameBuf
 		if _, err := io.ReadFull(br, sizeBuf); err != nil {
 			if err == io.EOF {
 				// End of stream
 				putUnmarshalWork(uw)
-				wg.Wait()
+				ctx.wg.Wait()
-				return callbackErr
+				return ctx.err
 			}
 			readErrors.Inc()
-			wg.Wait()
+			ctx.wg.Wait()
 			return fmt.Errorf("cannot read metricName size: %w", err)
 		}
 		readCalls.Inc()
 		bufSize := encoding.UnmarshalUint32(sizeBuf)
 		if bufSize > 1024*1024 {
 			parseErrors.Inc()
-			wg.Wait()
+			ctx.wg.Wait()
 			return fmt.Errorf("too big metricName size; got %d; shouldn't exceed %d", bufSize, 1024*1024)
 		}
 		uw.metricNameBuf = bytesutil.ResizeNoCopyMayOverallocate(uw.metricNameBuf, int(bufSize))
 		if _, err := io.ReadFull(br, uw.metricNameBuf); err != nil {
 			readErrors.Inc()
-			wg.Wait()
+			ctx.wg.Wait()
 			return fmt.Errorf("cannot read metricName with size %d bytes: %w", bufSize, err)
 		}
 		readCalls.Inc()
@@ -100,30 +80,36 @@ func ParseStream(r io.Reader, isGzip bool, callback func(block *Block) error) er
 		// Read uw.blockBuf
 		if _, err := io.ReadFull(br, sizeBuf); err != nil {
 			readErrors.Inc()
-			wg.Wait()
+			ctx.wg.Wait()
 			return fmt.Errorf("cannot read native block size: %w", err)
 		}
 		readCalls.Inc()
 		bufSize = encoding.UnmarshalUint32(sizeBuf)
 		if bufSize > 1024*1024 {
 			parseErrors.Inc()
-			wg.Wait()
+			ctx.wg.Wait()
 			return fmt.Errorf("too big native block size; got %d; shouldn't exceed %d", bufSize, 1024*1024)
 		}
 		uw.blockBuf = bytesutil.ResizeNoCopyMayOverallocate(uw.blockBuf, int(bufSize))
 		if _, err := io.ReadFull(br, uw.blockBuf); err != nil {
 			readErrors.Inc()
-			wg.Wait()
+			ctx.wg.Wait()
 			return fmt.Errorf("cannot read native block with size %d bytes: %w", bufSize, err)
 		}
 		readCalls.Inc()
 		blocksRead.Inc()
 
-		wg.Add(1)
+		ctx.wg.Add(1)
 		common.ScheduleUnmarshalWork(uw)
 	}
 }
 
+type streamContext struct {
+	wg      sync.WaitGroup
+	errLock sync.Mutex
+	err     error
+}
+
 // Block is a single block from `/api/v1/import/native` request.
 type Block struct {
 	MetricName storage.MetricName
@@ -149,13 +135,15 @@ var (
 
 type unmarshalWork struct {
 	tr            storage.TimeRange
-	callback      func(block *Block, parseErr error)
+	ctx           *streamContext
+	callback      func(block *Block) error
 	metricNameBuf []byte
 	blockBuf      []byte
 	block         Block
 }
 
 func (uw *unmarshalWork) reset() {
+	uw.ctx = nil
 	uw.callback = nil
 	uw.metricNameBuf = uw.metricNameBuf[:0]
 	uw.blockBuf = uw.blockBuf[:0]
@@ -167,8 +155,19 @@ func (uw *unmarshalWork) Unmarshal() {
 	err := uw.unmarshal()
 	if err != nil {
 		parseErrors.Inc()
+	} else {
+		err = uw.callback(&uw.block)
 	}
-	uw.callback(&uw.block, err)
+	ctx := uw.ctx
+	if err != nil {
+		processErrors.Inc()
+		ctx.errLock.Lock()
+		if ctx.err == nil {
+			ctx.err = fmt.Errorf("error when processing native block: %w", err)
+		}
+		ctx.errLock.Unlock()
+	}
+	ctx.wg.Done()
 	putUnmarshalWork(uw)
 }