Improve Any stack usage

[cdvr1993]

By doing this unsafe approach we are able to save one call to morestack.

The benchmark results are as follows:

goos: linux
goarch: amd64
pkg: go.uber.org/zap
cpu: AMD EPYC 7B13
BenchmarkAny/normal             19687398                61.96 ns/op
BenchmarkAny/normal-2           19890138                58.83 ns/op
BenchmarkAny/normal-4           20615840                58.00 ns/op
BenchmarkAny/normal-8           19891179                60.70 ns/op
BenchmarkAny/optimized          12379216                98.05 ns/op
BenchmarkAny/optimized-2        13090182                90.70 ns/op
BenchmarkAny/optimized-4        13424733                90.75 ns/op
BenchmarkAny/optimized-8        12764323                94.82 ns/op
BenchmarkAny/normal_with_logger                  7346602               154.2 ns/op
BenchmarkAny/normal_with_logger-2                8581615               145.1 ns/op
BenchmarkAny/normal_with_logger-4                8807738               137.8 ns/op
BenchmarkAny/normal_with_logger-8                8484144               146.7 ns/op
BenchmarkAny/optimized_with_logger               7519690               158.1 ns/op
BenchmarkAny/optimized_with_logger-2             7529078               154.9 ns/op
BenchmarkAny/optimized_with_logger-4             8006461               153.5 ns/op
BenchmarkAny/optimized_with_logger-8             7665606               150.2 ns/op
BenchmarkAny/normal_new_goroutine                 275497              3982 ns/op
BenchmarkAny/normal_new_goroutine-2               564081              2177 ns/op
BenchmarkAny/normal_new_goroutine-4               973549              1107 ns/op
BenchmarkAny/normal_new_goroutine-8              1856415               671.3 ns/op
BenchmarkAny/optimized_new_goroutine              550318              2073 ns/op
BenchmarkAny/optimized_new_goroutine-2           1000000              1275 ns/op
BenchmarkAny/optimized_new_goroutine-4           1690797               704.9 ns/op
BenchmarkAny/optimized_new_goroutine-8           1815640               642.1 ns/op
PASS
ok      go.uber.org/zap 34.181s

• The first set of benchmark where we only use runtime.KeepAlive are slower.
• The second set when we use logger the cost is the same.
• The third set is 1-2x faster because of stack growth.

The assembly shows the following change in stack usage:

# from
SUBQ $0x12f8, SP   // 4856 bytes

# to
SUBQ $0x3d8, SP  // 984 bytes

So this is ~5x reduction in stack usage.

Additionally I did a cpu profile, and there I found way more usage from the old approach on high core core (32 cores):

goos: linux
goarch: amd64
pkg: go.uber.org/zap
cpu: AMD EPYC 7B13
BenchmarkAny/normal_new_goroutine-32             4000000               664.5 ns/op
BenchmarkAny/optimized_new_goroutine-32          4000000               645.2 ns/op
PASS
ok      go.uber.org/zap 5.418s

Not sure why the time is the same but the profile isn't, only guess is that the goroutine scheduler is adding that latency?

Alternatives

Just fyi. I tried storing each switch statement in a function variable to avoid the compiler to allocate space for each function call, but by doing that the escape analyzer decides that the variables need to escape to the heap and we got the following results:

goos: linux
goarch: amd64
pkg: go.uber.org/zap
cpu: AMD EPYC 7B13
BenchmarkAny/normal             19453904                64.28 ns/op
BenchmarkAny/normal-2           19794348                60.95 ns/op
BenchmarkAny/normal-4           20221657                60.11 ns/op
BenchmarkAny/normal-8           19663383                61.51 ns/op
BenchmarkAny/optimized           6851958               176.9 ns/op
BenchmarkAny/optimized-2         7244982               165.2 ns/op
BenchmarkAny/optimized-4         7224718               160.4 ns/op
BenchmarkAny/optimized-8         6786295               171.4 ns/op
BenchmarkAny/normal_with_logger                  7592905               159.1 ns/op
BenchmarkAny/normal_with_logger-2                8350288               148.6 ns/op
BenchmarkAny/normal_with_logger-4                7812310               152.3 ns/op
BenchmarkAny/normal_with_logger-8                7676191               150.7 ns/op
BenchmarkAny/optimized_with_logger               5043325               239.7 ns/op
BenchmarkAny/optimized_with_logger-2             5618706               226.6 ns/op
BenchmarkAny/optimized_with_logger-4             5382121               215.1 ns/op
BenchmarkAny/optimized_with_logger-8             5612380               222.6 ns/op
BenchmarkAny/normal_new_goroutine                 262658              3889 ns/op
BenchmarkAny/normal_new_goroutine-2               649510              2127 ns/op
BenchmarkAny/normal_new_goroutine-4               869673              1229 ns/op
BenchmarkAny/normal_new_goroutine-8              1740733               699.0 ns/op
BenchmarkAny/optimized_new_goroutine              911757              1183 ns/op
BenchmarkAny/optimized_new_goroutine-2           1370961               929.5 ns/op
BenchmarkAny/optimized_new_goroutine-4           3625549               388.7 ns/op
BenchmarkAny/optimized_new_goroutine-8           2873199               389.7 ns/op
PASS
ok      go.uber.org/zap 37.317s

As you can see:

• optimized is almost 3x slower (compared to 50% before).
• optimized with logger is ~50% slower (compared to before when the cost was the same).
• optimized on a new goroutine is 2-3x faster (compared to before where it was only 1-2x).

Stack usage became:

SUBQ $0x58, SP // 88 bytes

This would be a reduction of >50x but with the cost of 2 more heap allocations.

[codecov]

Codecov Report

Merging #1301 (347f847) into master (24b7977) will increase coverage by 0.07%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##           master    #1301      +/-   ##
==========================================
+ Coverage   98.08%   98.15%   +0.07%     
==========================================
  Files          50       50              
  Lines        3242     3369     +127     
==========================================
+ Hits         3180     3307     +127     
  Misses         53       53              
  Partials        9        9              

Files Changed	Coverage Δ
field.go	100.00% <100.00%> (ø)

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[cdvr1993]

I believe the tradeoff here is not worth. While this approach's improvements on new goroutine stack allocation scenario is certainly impressive, the more frequent case is when we already have a stable goroutine whose stack size isn't growing anymore and we're logging happily. (the normal_with_logger case)

Let's discuss options in https://github.com/uber-go/zap/pull/1301.

@sywhang Is it really true?

Imagine you have 2 endpoints, one uses zap.Any the other doesn't. So one has a 2KB the other one 8KB stack. Imagine we have 75/25 traffic (25% to zap.Any). Then the average stack size in this case is going to be 4KB. So 25% of the traffic would need to allocate more stack.

[sywhang]

@cdvr1993 That's assuming each request is served by a fresh new goroutine though, right?

[cdvr1993]

@cdvr1993 That's assuming each request is served by a fresh new goroutine though, right?

Yes, the way some rpc frameworks work. For instance, gRPC:

https://github.com/grpc/grpc-go/blob/master/server.go#L996

func (s *Server) serveStreams(st transport.ServerTransport) {
	defer st.Close(errors.New("finished serving streams for the server transport"))
	var wg sync.WaitGroup

	st.HandleStreams(func(stream *transport.Stream) {
		wg.Add(1)
		if s.opts.numServerWorkers > 0 {
			data := &serverWorkerData{st: st, wg: &wg, stream: stream}
			select {
			case s.serverWorkerChannel <- data:
				return
			default:
				// If all stream workers are busy, fallback to the default code path.
			}
		}
		go func() {
			defer wg.Done()
			s.handleStream(st, stream, s.traceInfo(st, stream))
		}()
	}, func(ctx context.Context, method string) context.Context {
		if !EnableTracing {
			return ctx
		}
		tr := trace.New("grpc.Recv."+methodFamily(method), method)
		return trace.NewContext(ctx, tr)
	})
	wg.Wait()
}

[abhinav]

I know that @rabbbit's version is still WIP, but looking at b8b64dc, if that works, I think maybe that version may be more digestible. I'm not a fan of adding this much unsafe and //go:* by hand. If we end up needing to go to the unsafe route, we can discuss options around maintainability.

[cdvr1993]

I cleaned up a little bit the usage of go directives and unsafe usage:

cdvr1993@9b0f077

Stack is now 280 bytes vs 4856 bytes that we initially had.

Unsafe usage only remains for the interface parameter because we need a common type for it.

[prashantv]

I know that @rabbbit's version is still WIP, but looking at b8b64dc, if that works, I think maybe that version may be more digestible. I'm not a fan of adding this much unsafe and //go:* by hand. If we end up needing to go to the unsafe route, we can discuss options around maintainability.

Big +1, this is a huge amount of unsafe. If we end up needing this sort of unsafe usage for performance reasons, I don't think it should be in the default build, it should be under a separate build tag like fast_unsafe or something, so users can opt-in to the faster but unsafe version, while having the standard build use safe constructs.

[cdvr1993]

@prashantv @abhinav what do you think of creating a second Any function, like AnyUnsafe(), that way it is up to the users to use the unsafe version.

[abhinav]

First, let me say @cdvr1993, this is both, impressive debugging and nice workaround. I think we neglected to mention that in the initial reaction to the unsafe code.

The idea of "UnsafeAny" is more digestible, if a bit distasteful. I would still prefer to avoid it.

This is all a lot of debt to work around an (AFAIK) language issue. Please understand that part of the reason for my push back is that besides just the inherent risk of unsafe code, this makes it risky and difficult for anyone else to touch this code. The level of scrutiny necessary for every change here will have to be a lot higher, which is additional burden on maintainers.

If the unsafe option was the only way (and that's not yet settled), there are some mitigations that, in my opinion, would make the unsafe code more digestible:

• To address readability and maintainability, it would be preferable if the function and the accompanying helpers were all code-generated from a simpler, more readable source-of-truth. (What form that source of truth takes is up for discussion.)
• I think it would still be preferable to use build tags (//go:build !unsafe_zap and //go:build unsafe_zap) to select between the safe and unsafe implementations. My preference is unsafe-opt-in, but that's up for discussion if others feel differently.

Again, this is just my opinion. Other maintainers, feel free to chime in.
I don't recommend making the changes I just suggested until there's some agreement here, which I suspect is still contingent on the results of #1303 which avoids unsafe code entirely.

[cdvr1993]

Ok, just for the sake of showing how this can be cleaned up a bit, I created: #1304