Skip to content

Commit 8caa934

Browse files
kyasbalkyasbal
authored
add v6 khi file builder (#652)
1 parent 4fc3332 commit 8caa934

4 files changed

Lines changed: 437 additions & 0 deletions

File tree

‎pkg/model/khifile/v6/builder.go‎

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
// Copyright 2026 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package khifilev6
16+
17+
import (
18+
"fmt"
19+
"io"
20+
"iter"
21+
"slices"
22+
23+
pb "github.com/GoogleCloudPlatform/khi/pkg/generated/khifile/v6"
24+
"github.com/GoogleCloudPlatform/khi/pkg/model/khifile/v6/style"
25+
)
26+
27+
// Builder orchestrates the accumulators, pools, and final file generation for KHI v6 format.
28+
type Builder struct {
29+
idGenerator *IDGenerator
30+
internPool *InternPool
31+
TimelineAccumulator *TimelineAccumulator
32+
LogAccumulator *LogAccumulator
33+
MetadataAccumulator *MetadataAccumulator
34+
}
35+
36+
// NewBuilder initializes a new v6 Builder with all necessary accumulators and pools.
37+
func NewBuilder() *Builder {
38+
gen := &IDGenerator{}
39+
internPool := NewInternPool(gen)
40+
41+
return &Builder{
42+
idGenerator: gen,
43+
internPool: internPool,
44+
TimelineAccumulator: NewTimelineAccumulator(gen, internPool),
45+
LogAccumulator: NewLogAccumulator(internPool, gen),
46+
MetadataAccumulator: NewMetadataAccumulator(),
47+
}
48+
}
49+
50+
// Build writes the accumulated data to the provided io.Writer in KHI v6 format.
51+
func (b *Builder) Build(w io.Writer) error {
52+
writer, err := NewWriter(w)
53+
if err != nil {
54+
return fmt.Errorf("failed to create writer: %w", err)
55+
}
56+
57+
// 1. Write TimelineStyleChunk directly (no generator needed since it's a single chunk)
58+
styleChunk := style.GenerateChunk()
59+
if err := writer.WriteChunk(ChunkTypeTimelineStyle, styleChunk); err != nil {
60+
return fmt.Errorf("failed to write timeline style chunk: %w", err)
61+
}
62+
63+
// 2. Write MetadataChunk
64+
metadataList := b.MetadataAccumulator.Accumulate()
65+
if len(metadataList) > 0 {
66+
metadataGen := NewMetadataGenerator(slices.Values(metadataList))
67+
defer metadataGen.Close()
68+
if err := writer.WriteGenerator(metadataGen); err != nil {
69+
return fmt.Errorf("failed to write metadata chunk: %w", err)
70+
}
71+
}
72+
73+
// 3. Write LogChunks
74+
logs := b.LogAccumulator.Accumulate()
75+
if len(logs) > 0 {
76+
logGen := NewLogGenerator(slices.Values(logs))
77+
defer logGen.Close()
78+
if err := writer.WriteGenerator(logGen); err != nil {
79+
return fmt.Errorf("failed to write log chunks: %w", err)
80+
}
81+
}
82+
83+
// 4. Write TimelineChunks
84+
timelines, timelineItems := b.TimelineAccumulator.Accumulate()
85+
86+
if len(timelines) > 0 {
87+
timelineGen := NewTimelineGenerator(slices.Values(timelines))
88+
defer timelineGen.Close()
89+
if err := writer.WriteGenerator(timelineGen); err != nil {
90+
return fmt.Errorf("failed to write timeline chunks: %w", err)
91+
}
92+
}
93+
94+
if len(timelineItems) > 0 {
95+
timelineItemsGen := NewTimelineItemsGenerator(slices.Values(timelineItems))
96+
defer timelineItemsGen.Close()
97+
if err := writer.WriteGenerator(timelineItemsGen); err != nil {
98+
return fmt.Errorf("failed to write timeline items chunks: %w", err)
99+
}
100+
}
101+
102+
// 5. Write InternPoolChunk (Strings and FieldPathSets)
103+
stringSeq := mapSeq(b.internPool.SortedStringRefs(), func(ref *InternStringRef) *pb.InternString {
104+
return ref.ToProto()
105+
})
106+
stringGen := NewInternPoolGenerator(stringSeq)
107+
if err := writer.WriteGenerator(stringGen); err != nil {
108+
return fmt.Errorf("failed to write intern string chunks: %w", err)
109+
}
110+
111+
fieldSetSeq := mapSeq(b.internPool.FieldSetRefs(), func(ref *FieldPathSetRef) *pb.InternFieldPathSet {
112+
return ref.ToProto()
113+
})
114+
fieldPathSetGen := NewInternFieldPathSetGenerator(fieldSetSeq)
115+
if err := writer.WriteGenerator(fieldPathSetGen); err != nil {
116+
return fmt.Errorf("failed to write intern field path set chunks: %w", err)
117+
}
118+
119+
return nil
120+
}
121+
122+
func mapSeq[T any, U any](seq iter.Seq[T], f func(T) U) iter.Seq[U] {
123+
return func(yield func(U) bool) {
124+
for v := range seq {
125+
if !yield(f(v)) {
126+
return
127+
}
128+
}
129+
}
130+
}

‎pkg/model/khifile/v6/builder_test.go‎

Lines changed: 255 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,255 @@
1+
// Copyright 2026 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package khifilev6
16+
17+
import (
18+
"bytes"
19+
"errors"
20+
"io"
21+
"testing"
22+
"time"
23+
24+
"github.com/GoogleCloudPlatform/khi/pkg/common/structured"
25+
inspectionmetadata "github.com/GoogleCloudPlatform/khi/pkg/core/inspection/metadata"
26+
pb "github.com/GoogleCloudPlatform/khi/pkg/generated/khifile/v6"
27+
"github.com/GoogleCloudPlatform/khi/pkg/model/log"
28+
"github.com/google/go-cmp/cmp"
29+
"google.golang.org/protobuf/proto"
30+
"google.golang.org/protobuf/types/known/timestamppb"
31+
)
32+
33+
// TestBuilder_Build verifies that the Builder successfully generates a KHI file.
34+
// It tests both an empty build and a full build with metadata, logs, and timelines.
35+
func TestBuilder_Build(t *testing.T) {
36+
testCases := []struct {
37+
name string
38+
setup func(b *Builder)
39+
verify func(t *testing.T, reader *Reader)
40+
}{
41+
{
42+
name: "empty build writes only style chunk",
43+
setup: func(b *Builder) {},
44+
verify: func(t *testing.T, reader *Reader) {
45+
chunk, err := reader.NextChunk()
46+
if err != nil {
47+
t.Fatalf("Failed to read next chunk: %v.", err)
48+
}
49+
if chunk.Type != ChunkTypeTimelineStyle {
50+
t.Errorf("Expected style chunk, got %v.", chunk.Type)
51+
}
52+
53+
_, err = reader.NextChunk()
54+
if !errors.Is(err, io.EOF) {
55+
t.Errorf("Expected EOF, got %v.", err)
56+
}
57+
},
58+
},
59+
{
60+
name: "full build writes metadata, log, timeline, style, and intern chunks",
61+
setup: func(b *Builder) {
62+
_ = b.MetadataAccumulator.AddMetadata(&inspectionmetadata.HeaderMetadata{
63+
InspectionType: "test-inspection-type",
64+
InspectionName: "test-inspection-name",
65+
FileSize: 1024,
66+
})
67+
68+
node := structured.NewStandardMap(
69+
[]string{"message"},
70+
[]structured.Node{
71+
structured.NewStandardScalarNode("hello log"),
72+
},
73+
)
74+
severityID := uint32(1)
75+
logTypeID := uint32(2)
76+
_ = b.LogAccumulator.AddLog(&StagingLog{
77+
Log: log.NewLog(structured.NewNodeReader(node)),
78+
Summary: "hello summary",
79+
Timestamp: time.Date(2026, 4, 29, 8, 0, 0, 0, time.UTC),
80+
Severity: &pb.Severity{Id: &severityID},
81+
LogType: &pb.LogType{Id: &logTypeID},
82+
})
83+
84+
timelineTypeID := uint32(3)
85+
timelineType := &pb.TimelineType{Id: &timelineTypeID}
86+
path := b.TimelineAccumulator.GetPath(nil, PathSegment{
87+
Name: "test-timeline-path",
88+
Type: timelineType,
89+
})
90+
tb := b.TimelineAccumulator.GetBuilder(path)
91+
logID := uint32(1)
92+
tb.AddRevision(&pb.Revision{
93+
LogId: &logID,
94+
ChangedTime: timestamppb.New(time.Date(2026, 4, 29, 8, 0, 0, 0, time.UTC)),
95+
})
96+
},
97+
verify: func(t *testing.T, reader *Reader) {
98+
var chunks []*Chunk
99+
for {
100+
chunk, err := reader.NextChunk()
101+
if errors.Is(err, io.EOF) {
102+
break
103+
}
104+
if err != nil {
105+
t.Fatalf("Failed to read next chunk: %v.", err)
106+
}
107+
chunks = append(chunks, chunk)
108+
}
109+
110+
expectedTypes := []ChunkType{
111+
ChunkTypeTimelineStyle,
112+
ChunkTypeMetadata,
113+
ChunkTypeLog,
114+
ChunkTypeTimeline,
115+
ChunkTypeTimeline,
116+
ChunkTypeInternPool,
117+
ChunkTypeInternPool,
118+
}
119+
120+
var gotTypes []ChunkType
121+
for _, c := range chunks {
122+
gotTypes = append(gotTypes, c.Type)
123+
}
124+
125+
if diff := cmp.Diff(expectedTypes, gotTypes); diff != "" {
126+
t.Errorf("Chunk types mismatch (-want +got):\n%s", diff)
127+
}
128+
129+
for _, c := range chunks {
130+
switch c.Type {
131+
case ChunkTypeTimelineStyle:
132+
var styleChunk pb.TimelineStyleChunk
133+
if err := proto.Unmarshal(c.Data, &styleChunk); err != nil {
134+
t.Fatalf("Failed to unmarshal style chunk: %v.", err)
135+
}
136+
if styleChunk.IconAtlas == nil {
137+
t.Error("Expected style chunk to have non-nil IconAtlas.")
138+
}
139+
140+
case ChunkTypeMetadata:
141+
var metadataChunk pb.MetadataChunk
142+
if err := proto.Unmarshal(c.Data, &metadataChunk); err != nil {
143+
t.Fatalf("Failed to unmarshal metadata chunk: %v.", err)
144+
}
145+
if len(metadataChunk.Metadata) != 1 {
146+
t.Errorf("Expected 1 metadata item, got %d.", len(metadataChunk.Metadata))
147+
break
148+
}
149+
header := metadataChunk.Metadata[0].GetHeader()
150+
if header == nil {
151+
t.Error("Expected header metadata.")
152+
break
153+
}
154+
if header.GetInspectionType() != "test-inspection-type" || header.GetInspectionName() != "test-inspection-name" {
155+
t.Errorf("Unexpected header: %+v.", header)
156+
}
157+
158+
case ChunkTypeLog:
159+
var logChunk pb.LogChunk
160+
if err := proto.Unmarshal(c.Data, &logChunk); err != nil {
161+
t.Fatalf("Failed to unmarshal log chunk: %v.", err)
162+
}
163+
if len(logChunk.Logs) != 1 {
164+
t.Errorf("Expected 1 log, got %d.", len(logChunk.Logs))
165+
break
166+
}
167+
logItem := logChunk.Logs[0]
168+
if logItem.GetId() != 1 {
169+
t.Errorf("Expected log ID 1, got %d.", logItem.GetId())
170+
}
171+
if logItem.GetSeverityTypeId() != 1 {
172+
t.Errorf("Expected severity ID 1, got %d.", logItem.GetSeverityTypeId())
173+
}
174+
if logItem.GetLogTypeId() != 2 {
175+
t.Errorf("Expected log type ID 2, got %d.", logItem.GetLogTypeId())
176+
}
177+
178+
case ChunkTypeTimeline:
179+
var timelineChunk pb.TimelineChunk
180+
if err := proto.Unmarshal(c.Data, &timelineChunk); err != nil {
181+
t.Fatalf("Failed to unmarshal timeline chunk: %v.", err)
182+
}
183+
switch {
184+
case len(timelineChunk.Timelines) > 0:
185+
if len(timelineChunk.Timelines) != 1 {
186+
t.Errorf("Expected 1 timeline, got %d.", len(timelineChunk.Timelines))
187+
} else {
188+
tl := timelineChunk.Timelines[0]
189+
if tl.GetId() != 1 {
190+
t.Errorf("Expected timeline ID 1, got %d.", tl.GetId())
191+
}
192+
}
193+
case len(timelineChunk.TimelineItems) > 0:
194+
if len(timelineChunk.TimelineItems) != 1 {
195+
t.Errorf("Expected 1 timeline item collection, got %d.", len(timelineChunk.TimelineItems))
196+
} else {
197+
ti := timelineChunk.TimelineItems[0]
198+
if ti.GetId() != 1 {
199+
t.Errorf("Expected timeline items ID 1, got %d.", ti.GetId())
200+
}
201+
if len(ti.Revisions) != 1 {
202+
t.Errorf("Expected 1 revision, got %d.", len(ti.Revisions))
203+
} else {
204+
rev := ti.Revisions[0]
205+
if rev.GetLogId() != 1 {
206+
t.Errorf("Expected log ID 1 in revision, got %d.", rev.GetLogId())
207+
}
208+
}
209+
}
210+
default:
211+
t.Error("Expected either Timelines or TimelineItems in TimelineChunk.")
212+
}
213+
214+
case ChunkTypeInternPool:
215+
var internPool pb.InterningPoolChunk
216+
if err := proto.Unmarshal(c.Data, &internPool); err != nil {
217+
t.Fatalf("Failed to unmarshal intern pool chunk: %v.", err)
218+
}
219+
if len(internPool.Strings) > 0 {
220+
foundSummary := false
221+
for _, s := range internPool.Strings {
222+
if s.GetValue() == "hello summary" {
223+
foundSummary = true
224+
break
225+
}
226+
}
227+
if !foundSummary {
228+
t.Error("Expected 'hello summary' to be interned in InternPool.")
229+
}
230+
}
231+
}
232+
}
233+
},
234+
},
235+
}
236+
237+
for _, tc := range testCases {
238+
t.Run(tc.name, func(t *testing.T) {
239+
b := NewBuilder()
240+
tc.setup(b)
241+
242+
var buf bytes.Buffer
243+
if err := b.Build(&buf); err != nil {
244+
t.Fatalf("Build() failed: %v.", err)
245+
}
246+
247+
reader, err := NewReader(&buf)
248+
if err != nil {
249+
t.Fatalf("NewReader() failed: %v.", err)
250+
}
251+
252+
tc.verify(t, reader)
253+
})
254+
}
255+
}

‎pkg/model/khifile/v6/generator.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,9 @@ func (g *splittingGenerator[T, C]) Close() error {
128128
return nil
129129
}
130130

131+
// splittingGenerator implements ChunkGenerator
132+
var _ ChunkGenerator = (*splittingGenerator[proto.Message, proto.Message])(nil)
133+
131134
// NewInternPoolGenerator creates a SplittingGenerator for InternPool chunks.
132135
// It groups pb.InternString messages into pb.InterningPoolChunk respecting the size limit.
133136
func NewInternPoolGenerator(seq iter.Seq[*pb.InternString]) ChunkGenerator {

0 commit comments

Comments
 (0)