chore(go.d.plugin): improve function parser (#19143)
Ilya Mashchenko committed
Dec 6, 2024 at 16:08 UTC
a1d0fe557f93479c1acfd030f1f014348b02e1c2
3 files changed
+134
-113
src/go/plugin/go.d/agent/functions/function.go
deleted
-96
@@ -1,96 +0,0 @@
1
-// SPDX-License-Identifier: GPL-3.0-or-later
2
-
3
-package functions
4
-
5
-import (
6
- "bytes"
7
- "context"
8
- "encoding/csv"
9
- "fmt"
10
- "strconv"
11
- "strings"
12
- "time"
13
-)
14
-
15
-type Function struct {
16
- key string
17
- UID string
18
- Timeout time.Duration
19
- Name string
20
- Args []string
21
- Payload []byte
22
- Permissions string
23
- Source string
24
- ContentType string
25
-}
26
-
27
-func (f *Function) String() string {
28
- return fmt.Sprintf("key: '%s', uid: '%s', timeout: '%s', function: '%s', args: '%v', permissions: '%s', source: '%s', contentType: '%s', payload: '%s'",
29
- f.key, f.UID, f.Timeout, f.Name, f.Args, f.Permissions, f.Source, f.ContentType, string(f.Payload))
30
-}
31
-
32
-func parseFunction(s string) (*Function, error) {
33
- r := csv.NewReader(strings.NewReader(s))
34
- r.Comma = ' '
35
-
36
- parts, err := r.Read()
37
- if err != nil {
38
- return nil, err
39
- }
40
-
41
- // FUNCTION UID Timeout "Name ...Parameters" 0xPermissions "SourceType" [ContentType]
42
- if n := len(parts); n != 6 && n != 7 {
43
- return nil, fmt.Errorf("unexpected number of words: want 6 or 7, got %d (%v)", n, parts)
44
- }
45
-
46
- timeout, err := strconv.ParseInt(parts[2], 10, 64)
47
- if err != nil {
48
- return nil, err
49
- }
50
-
51
- cmd := strings.Split(parts[3], " ")
52
-
53
- fn := &Function{
54
- key: parts[0],
55
- UID: parts[1],
56
- Timeout: time.Duration(timeout) * time.Second,
57
- Name: cmd[0],
58
- Args: cmd[1:],
59
- Permissions: parts[4],
60
- Source: parts[5],
61
- }
62
-
63
- if len(parts) == 7 {
64
- fn.ContentType = parts[6]
65
- }
66
-
67
- return fn, nil
68
-}
69
-
70
-func parseFunctionWithPayload(ctx context.Context, s string, in input) (*Function, error) {
71
- fn, err := parseFunction(s)
72
- if err != nil {
73
- return nil, err
74
- }
75
-
76
- var buf bytes.Buffer
77
-
78
- for {
79
- select {
80
- case <-ctx.Done():
81
- return nil, nil
82
- case line, ok := <-in.lines():
83
- if !ok {
84
- return nil, nil
85
- }
86
- if line == "FUNCTION_PAYLOAD_END" {
87
- fn.Payload = append(fn.Payload, buf.Bytes()...)
88
- return fn, nil
89
- }
90
- if buf.Len() > 0 {
91
- buf.WriteString("\n")
92
- }
93
- buf.WriteString(line)
94
- }
95
- }
96
-}
src/go/plugin/go.d/agent/functions/manager.go
+4
-17
@@ -8,7 +8,6 @@ import (
8
"fmt"
9
"log/slog"
10
"strconv"
11
- "strings"
11
"sync"
12
"time"
13
@@ -55,6 +54,8 @@ func (m *Manager) Run(ctx context.Context, quitCh chan struct{}) {
54
}
55
56
func (m *Manager) run(ctx context.Context, quitCh chan struct{}) {
57
+ parser := newInputParser()
58
+
59
for {
60
select {
61
case <-ctx.Done():
@@ -63,29 +64,15 @@ func (m *Manager) run(ctx context.Context, quitCh chan struct{}) {
64
if !ok {
65
return
66
}
66
-
67
- var fn *Function
68
- var err error
69
-
70
- // FIXME: if we are waiting for FUNCTION_PAYLOAD_END and a new FUNCTION* appears,
71
- // we need to discard the current one and switch to the new one
72
- switch {
73
- case strings.HasPrefix(line, "FUNCTION "):
74
- fn, err = parseFunction(line)
75
- case strings.HasPrefix(line, "FUNCTION_PAYLOAD "):
76
- fn, err = parseFunctionWithPayload(ctx, line, m.input)
77
- case line == "":
78
- continue
79
- case line == "QUIT":
67
+ if line == "QUIT" {
68
if quitCh != nil {
69
quitCh <- struct{}{}
70
return
71
}
84
- default:
85
- m.Warningf("unexpected line: '%s'", line)
72
continue
73
}
74
75
+ fn, err := parser.parse(line)
76
if err != nil {
77
m.Warningf("parse function: %v ('%s')", err, line)
78
continue
src/go/plugin/go.d/agent/functions/parser.go
new
+130
@@ -0,0 +1,130 @@
1
+// SPDX-License-Identifier: GPL-3.0-or-later
2
+
3
+package functions
4
+
5
+import (
6
+ "bytes"
7
+ "encoding/csv"
8
+ "errors"
9
+ "fmt"
10
+ "strconv"
11
+ "strings"
12
+ "time"
13
+)
14
+
15
+type Function struct {
16
+ key string
17
+ UID string
18
+ Timeout time.Duration
19
+ Name string
20
+ Args []string
21
+ Payload []byte
22
+ Permissions string
23
+ Source string
24
+ ContentType string
25
+}
26
+
27
+func (f *Function) String() string {
28
+ return fmt.Sprintf("key: '%s', uid: '%s', timeout: '%s', function: '%s', args: '%v', permissions: '%s', source: '%s', contentType: '%s', payload: '%s'",
29
+ f.key, f.UID, f.Timeout, f.Name, f.Args, f.Permissions, f.Source, f.ContentType, string(f.Payload))
30
+}
31
+
32
+func newInputParser() *inputParser {
33
+ return &inputParser{}
34
+}
35
+
36
+type inputParser struct {
37
+ currentFn *Function
38
+ readingPayload bool
39
+ payloadBuf bytes.Buffer
40
+}
41
+
42
+func (p *inputParser) parse(line string) (*Function, error) {
43
+ if line = strings.TrimSpace(line); line == "" {
44
+ return nil, nil
45
+ }
46
+
47
+ if p.readingPayload {
48
+ return p.handlePayloadLine(line)
49
+ }
50
+
51
+ switch {
52
+ case strings.HasPrefix(line, "FUNCTION "):
53
+ return p.parseFunction(line)
54
+ case strings.HasPrefix(line, "FUNCTION_PAYLOAD "):
55
+ fn, err := p.parseFunction(line)
56
+ if err != nil {
57
+ return nil, err
58
+ }
59
+ p.readingPayload = true
60
+ p.currentFn = fn
61
+ p.payloadBuf.Reset()
62
+ return nil, nil
63
+ default:
64
+ return nil, errors.New("unexpected line format")
65
+ }
66
+}
67
+
68
+func (p *inputParser) handlePayloadLine(line string) (*Function, error) {
69
+ if line == "FUNCTION_PAYLOAD_END" {
70
+ p.readingPayload = false
71
+ p.currentFn.Payload = []byte(p.payloadBuf.String())
72
+ fn := p.currentFn
73
+ p.currentFn = nil
74
+ return fn, nil
75
+ }
76
+
77
+ if strings.HasPrefix(line, "FUNCTION") {
78
+ p.readingPayload = false
79
+ p.currentFn = nil
80
+ p.payloadBuf.Reset()
81
+ return p.parse(line)
82
+ }
83
+
84
+ if p.payloadBuf.Len() > 0 {
85
+ p.payloadBuf.WriteByte('\n')
86
+ }
87
+ p.payloadBuf.WriteString(line)
88
+
89
+ return nil, nil
90
+}
91
+
92
+func (p *inputParser) parseFunction(line string) (*Function, error) {
93
+ r := csv.NewReader(strings.NewReader(line))
94
+ r.Comma = ' '
95
+
96
+ parts, err := r.Read()
97
+ if err != nil {
98
+ return nil, fmt.Errorf("failed to parse CSV: %w", err)
99
+ }
100
+
101
+ if n := len(parts); n != 6 && n != 7 {
102
+ return nil, fmt.Errorf("unexpected number of parts: want 6 or 7, got %d", n)
103
+ }
104
+
105
+ timeout, err := strconv.ParseInt(parts[2], 10, 64)
106
+ if err != nil {
107
+ return nil, fmt.Errorf("invalid timeout value: %w", err)
108
+ }
109
+
110
+ nameAndArgs := strings.Split(parts[3], " ")
111
+ if len(nameAndArgs) == 0 {
112
+ return nil, fmt.Errorf("empty function name and arguments")
113
+ }
114
+
115
+ fn := &Function{
116
+ key: parts[0],
117
+ UID: parts[1],
118
+ Timeout: time.Duration(timeout) * time.Second,
119
+ Name: nameAndArgs[0],
120
+ Args: nameAndArgs[1:],
121
+ Permissions: parts[4],
122
+ Source: parts[5],
123
+ }
124
+
125
+ if len(parts) == 7 {
126
+ fn.ContentType = parts[6]
127
+ }
128
+
129
+ return fn, nil
130
+}