Add Unit to create_text and parse_text

Signed-off-by: Arianna Vespri <[email protected]>

Author: vesari

expfmt/text_create.go

@@ -148,6 +148,33 @@ func MetricFamilyToText(out io.Writer, in *dto.MetricFamily) (written int, err e
 	if err != nil {
 		return
 	}
+	if in.Unit != nil {
+		n, err = w.WriteString("# UNIT ")
+		written += n
+		if err != nil {
+			return
+		}
+		n, err = w.WriteString(name)
+		written += n
+		if err != nil {
+			return
+		}
+		err = w.WriteByte(' ')
+		written++
+		if err != nil {
+			return
+		}
+		n, err = writeEscapedString(w, *in.Unit, false)
+		written += n
+		if err != nil {
+			return
+		}
+		err = w.WriteByte('\n')
+		written++
+		if err != nil {
+			return
+		}
+	}
 
 	// Finally the samples, one line for each.
 	for _, metric := range in.Metric {

expfmt/text_parse.go

@@ -178,7 +178,7 @@ func (p *TextParser) startComment() stateFn {
 		return p.startOfLine
 	}
 	keyword := p.currentToken.String()
-	if keyword != "HELP" && keyword != "TYPE" {
+	if keyword != "HELP" && keyword != "TYPE" && keyword != "UNIT" {
 		// Generic comment, ignore by fast forwarding to end of line.
 		for p.currentByte != '\n' {
 			if p.currentByte, p.err = p.buf.ReadByte(); p.err != nil {
@@ -217,6 +217,8 @@ func (p *TextParser) startComment() stateFn {
 		return p.readingHelp
 	case "TYPE":
 		return p.readingType
+	case "UNIT":
+		return p.readingUnit
 	}
 	panic(fmt.Sprintf("code error: unexpected keyword %q", keyword))
 }
@@ -528,6 +530,21 @@ func (p *TextParser) readingType() stateFn {
 	return p.startOfLine
 }
 
+// readingUnit represents the state where the last byte read (now in
+// p.currentByte) is the first byte of the docstring after 'UNIT'.
+func (p *TextParser) readingUnit() stateFn {
+	if p.currentMF.Unit != nil {
+		p.parseError(fmt.Sprintf("second UNIT line for metric name %q", p.currentMF.GetName()))
+		return nil
+	}
+	// Rest of line is the docstring.
+	if p.readTokenUntilNewline(true); p.err != nil {
+		return nil // Unexpected end of input.
+	}
+	p.currentMF.Unit = proto.String(p.currentToken.String())
+	return p.startOfLine
+}
+
 // parseError sets p.err to a ParseError at the current line with the given
 // message.
 func (p *TextParser) parseError(msg string) {