Skip to content

Commit 24bb4b1

Browse files
integrate latest wazero experimental updates (#66)
Pulls the changes from wazero/wazero#1453 --------- Signed-off-by: Achille Roussel <[email protected]>
1 parent c958a5e commit 24bb4b1

File tree

6 files changed

+55
-42
lines changed

6 files changed

+55
-42
lines changed

cpu.go

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -190,7 +190,7 @@ func (p *CPUProfiler) NewFunctionListener(def api.FunctionDefinition) experiment
190190

191191
type cpuProfiler struct{ *CPUProfiler }
192192

193-
func (p cpuProfiler) Before(ctx context.Context, mod api.Module, def api.FunctionDefinition, params []uint64, si experimental.StackIterator) context.Context {
193+
func (p cpuProfiler) Before(ctx context.Context, mod api.Module, def api.FunctionDefinition, _ []uint64, si experimental.StackIterator) {
194194
var frame cpuTimeFrame
195195
p.mutex.Lock()
196196

@@ -212,10 +212,9 @@ func (p cpuProfiler) Before(ctx context.Context, mod api.Module, def api.Functio
212212

213213
p.mutex.Unlock()
214214
p.frames = append(p.frames, frame)
215-
return ctx
216215
}
217216

218-
func (p cpuProfiler) After(ctx context.Context, mod api.Module, def api.FunctionDefinition, err error, results []uint64) {
217+
func (p cpuProfiler) After(ctx context.Context, mod api.Module, def api.FunctionDefinition, _ []uint64) {
219218
i := len(p.frames) - 1
220219
f := p.frames[i]
221220
p.frames = p.frames[:i]
@@ -229,3 +228,7 @@ func (p cpuProfiler) After(ctx context.Context, mod api.Module, def api.Function
229228
p.traces = append(p.traces, f.trace)
230229
}
231230
}
231+
232+
func (p cpuProfiler) Abort(ctx context.Context, mod api.Module, def api.FunctionDefinition, _ error) {
233+
p.After(ctx, mod, def, nil)
234+
}

go.mod

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,6 @@ go 1.20
44

55
require (
66
github.com/google/pprof v0.0.0-20230406165453-00490a63f317
7-
github.com/tetratelabs/wazero v1.1.1-0.20230509203401-ef3d67119550
7+
github.com/tetratelabs/wazero v1.1.1-0.20230511035210-78c35acd6e1c
88
golang.org/x/exp v0.0.0-20230425010034-47ecfdc1ba53
99
)

go.sum

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
github.com/google/pprof v0.0.0-20230406165453-00490a63f317 h1:hFhpt7CTmR3DX+b4R19ydQFtofxT0Sv3QsKNMVQYTMQ=
22
github.com/google/pprof v0.0.0-20230406165453-00490a63f317/go.mod h1:79YE0hCXdHag9sBkw2o+N/YnZtTkXi0UT9Nnixa5eYk=
3-
github.com/tetratelabs/wazero v1.1.1-0.20230509203401-ef3d67119550 h1:t2+P0h4kAojvOntox6816oFO76+N+lLYq2b3xKMzNnI=
4-
github.com/tetratelabs/wazero v1.1.1-0.20230509203401-ef3d67119550/go.mod h1:wYx2gNRg8/WihJfSDxA1TIL8H+GkfLYm+bIfbblu9VQ=
3+
github.com/tetratelabs/wazero v1.1.1-0.20230511035210-78c35acd6e1c h1:3vKw5AyUv+16qm9tIVfvm205aDIpDmrjxPydJ87GBKY=
4+
github.com/tetratelabs/wazero v1.1.1-0.20230511035210-78c35acd6e1c/go.mod h1:wYx2gNRg8/WihJfSDxA1TIL8H+GkfLYm+bIfbblu9VQ=
55
golang.org/x/exp v0.0.0-20230425010034-47ecfdc1ba53 h1:5llv2sWeaMSnA3w2kS57ouQQ4pudlXrR0dCgw51QK9o=
66
golang.org/x/exp v0.0.0-20230425010034-47ecfdc1ba53/go.mod h1:V1LtkGg67GoY2N1AnLN78QLrzxkLyJw7RJb1gzOOz9w=

mem.go

Lines changed: 33 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -221,16 +221,16 @@ type mallocProfiler struct {
221221
stack stackTrace
222222
}
223223

224-
func (p *mallocProfiler) Before(ctx context.Context, mod api.Module, def api.FunctionDefinition, params []uint64, si experimental.StackIterator) context.Context {
224+
func (p *mallocProfiler) Before(ctx context.Context, mod api.Module, def api.FunctionDefinition, params []uint64, si experimental.StackIterator) {
225225
p.size = api.DecodeU32(params[0])
226226
p.stack = makeStackTrace(p.stack, si)
227-
return ctx
228227
}
229228

230-
func (p *mallocProfiler) After(ctx context.Context, mod api.Module, def api.FunctionDefinition, err error, results []uint64) {
231-
if err == nil {
232-
p.memory.observeAlloc(api.DecodeU32(results[0]), p.size, p.stack)
233-
}
229+
func (p *mallocProfiler) After(ctx context.Context, mod api.Module, def api.FunctionDefinition, results []uint64) {
230+
p.memory.observeAlloc(api.DecodeU32(results[0]), p.size, p.stack)
231+
}
232+
233+
func (p *mallocProfiler) Abort(ctx context.Context, mod api.Module, def api.FunctionDefinition, _ error) {
234234
}
235235

236236
type callocProfiler struct {
@@ -240,17 +240,17 @@ type callocProfiler struct {
240240
stack stackTrace
241241
}
242242

243-
func (p *callocProfiler) Before(ctx context.Context, mod api.Module, def api.FunctionDefinition, params []uint64, si experimental.StackIterator) context.Context {
243+
func (p *callocProfiler) Before(ctx context.Context, mod api.Module, def api.FunctionDefinition, params []uint64, si experimental.StackIterator) {
244244
p.count = api.DecodeU32(params[0])
245245
p.size = api.DecodeU32(params[1])
246246
p.stack = makeStackTrace(p.stack, si)
247-
return ctx
248247
}
249248

250-
func (p *callocProfiler) After(ctx context.Context, mod api.Module, def api.FunctionDefinition, err error, results []uint64) {
251-
if err == nil {
252-
p.memory.observeAlloc(api.DecodeU32(results[0]), p.count*p.size, p.stack)
253-
}
249+
func (p *callocProfiler) After(ctx context.Context, mod api.Module, def api.FunctionDefinition, results []uint64) {
250+
p.memory.observeAlloc(api.DecodeU32(results[0]), p.count*p.size, p.stack)
251+
}
252+
253+
func (p *callocProfiler) Abort(ctx context.Context, mod api.Module, def api.FunctionDefinition, _ error) {
254254
}
255255

256256
type reallocProfiler struct {
@@ -260,34 +260,35 @@ type reallocProfiler struct {
260260
stack stackTrace
261261
}
262262

263-
func (p *reallocProfiler) Before(ctx context.Context, mod api.Module, def api.FunctionDefinition, params []uint64, si experimental.StackIterator) context.Context {
263+
func (p *reallocProfiler) Before(ctx context.Context, mod api.Module, def api.FunctionDefinition, params []uint64, si experimental.StackIterator) {
264264
p.addr = api.DecodeU32(params[0])
265265
p.size = api.DecodeU32(params[1])
266266
p.stack = makeStackTrace(p.stack, si)
267-
return ctx
268267
}
269268

270-
func (p *reallocProfiler) After(ctx context.Context, mod api.Module, def api.FunctionDefinition, err error, results []uint64) {
271-
if err == nil {
272-
p.memory.observeFree(p.addr)
273-
p.memory.observeAlloc(api.DecodeU32(results[0]), p.size, p.stack)
274-
}
269+
func (p *reallocProfiler) After(ctx context.Context, mod api.Module, def api.FunctionDefinition, results []uint64) {
270+
p.memory.observeFree(p.addr)
271+
p.memory.observeAlloc(api.DecodeU32(results[0]), p.size, p.stack)
272+
}
273+
274+
func (p *reallocProfiler) Abort(ctx context.Context, mod api.Module, def api.FunctionDefinition, _ error) {
275275
}
276276

277277
type freeProfiler struct {
278278
memory *MemoryProfiler
279279
addr uint32
280280
}
281281

282-
func (p *freeProfiler) Before(ctx context.Context, mod api.Module, def api.FunctionDefinition, params []uint64, si experimental.StackIterator) context.Context {
282+
func (p *freeProfiler) Before(ctx context.Context, mod api.Module, def api.FunctionDefinition, params []uint64, si experimental.StackIterator) {
283283
p.addr = api.DecodeU32(params[0])
284-
return ctx
285284
}
286285

287-
func (p *freeProfiler) After(ctx context.Context, mod api.Module, def api.FunctionDefinition, err error, results []uint64) {
288-
if err == nil {
289-
p.memory.observeFree(p.addr)
290-
}
286+
func (p *freeProfiler) After(ctx context.Context, mod api.Module, def api.FunctionDefinition, _ []uint64) {
287+
p.memory.observeFree(p.addr)
288+
}
289+
290+
func (p *freeProfiler) Abort(ctx context.Context, mod api.Module, def api.FunctionDefinition, _ error) {
291+
p.After(ctx, mod, def, nil)
291292
}
292293

293294
type goRuntimeMallocgcProfiler struct {
@@ -296,7 +297,7 @@ type goRuntimeMallocgcProfiler struct {
296297
stack stackTrace
297298
}
298299

299-
func (p *goRuntimeMallocgcProfiler) Before(ctx context.Context, mod api.Module, def api.FunctionDefinition, params []uint64, si experimental.StackIterator) context.Context {
300+
func (p *goRuntimeMallocgcProfiler) Before(ctx context.Context, mod api.Module, def api.FunctionDefinition, params []uint64, si experimental.StackIterator) {
300301
imod := mod.(experimental.InternalModule)
301302
mem := imod.Memory()
302303

@@ -309,13 +310,16 @@ func (p *goRuntimeMallocgcProfiler) Before(ctx context.Context, mod api.Module,
309310
} else {
310311
p.size = 0
311312
}
312-
return ctx
313313
}
314314

315-
func (p *goRuntimeMallocgcProfiler) After(ctx context.Context, mod api.Module, def api.FunctionDefinition, err error, results []uint64) {
316-
if err == nil && p.size != 0 {
315+
func (p *goRuntimeMallocgcProfiler) After(ctx context.Context, mod api.Module, def api.FunctionDefinition, _ []uint64) {
316+
if p.size != 0 {
317317
// TODO: get the returned pointer
318318
addr := uint32(0)
319319
p.memory.observeAlloc(addr, p.size, p.stack)
320320
}
321321
}
322+
323+
func (p *goRuntimeMallocgcProfiler) Abort(ctx context.Context, mod api.Module, def api.FunctionDefinition, _ error) {
324+
p.After(ctx, mod, def, nil)
325+
}

sampler.go

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -55,22 +55,27 @@ type sampledFunctionListener struct {
5555
lstn experimental.FunctionListener
5656
}
5757

58-
func (s *sampledFunctionListener) Before(ctx context.Context, mod api.Module, def api.FunctionDefinition, params []uint64, stack experimental.StackIterator) context.Context {
58+
func (s *sampledFunctionListener) Before(ctx context.Context, mod api.Module, def api.FunctionDefinition, params []uint64, stack experimental.StackIterator) {
5959
bit := uint(0)
6060

6161
if s.count--; s.count == 0 {
6262
s.count = s.cycle
63+
s.lstn.Before(ctx, mod, def, params, stack)
6364
bit = 1
64-
ctx = s.lstn.Before(ctx, mod, def, params, stack)
6565
}
6666

6767
s.stack.push(bit)
68-
return ctx
6968
}
7069

71-
func (s *sampledFunctionListener) After(ctx context.Context, mod api.Module, def api.FunctionDefinition, err error, results []uint64) {
70+
func (s *sampledFunctionListener) After(ctx context.Context, mod api.Module, def api.FunctionDefinition, results []uint64) {
7271
if s.stack.pop() != 0 {
73-
s.lstn.After(ctx, mod, def, err, results)
72+
s.lstn.After(ctx, mod, def, results)
73+
}
74+
}
75+
76+
func (s *sampledFunctionListener) Abort(ctx context.Context, mod api.Module, def api.FunctionDefinition, err error) {
77+
if s.stack.pop() != 0 {
78+
s.lstn.Abort(ctx, mod, def, err)
7479
}
7580
}
7681

sampler_test.go

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,10 +25,11 @@ func TestSampledFunctionListener(t *testing.T) {
2525

2626
function := module.Function(0).Definition()
2727
listener := factory.NewFunctionListener(function)
28+
ctx := context.Background()
2829

2930
for i := 0; i < 20; i++ {
30-
ctx := listener.Before(context.Background(), module, function, nil, nil)
31-
listener.After(ctx, module, function, nil, nil)
31+
listener.Before(ctx, module, function, nil, nil)
32+
listener.After(ctx, module, function, nil)
3233
}
3334

3435
if n != 2 {

0 commit comments

Comments
 (0)