Skip to content

Commit d5db6cf

Browse files
ivanauthivanauth
committed
fix: re-enable async logging with proper shutdown flushing
Fixes #889 Async logging was disabled in October 2022 (PR #888) because the diode.Writer buffer wasn't being flushed on shutdown, causing log messages to be lost. The cobrazerolog library doesn't expose the diode writer closer. This change: - Adds thread-safe global closer for diode.Writer in internal/logging - Creates ConfigureLoggingRunE() to replace cobrazerolog (exposes closer) - Restructures main() to run() int pattern with defer log.Close() - Converts log.Fatal() calls to return errors for proper cleanup - Fixes TTY detection to check stderr instead of stdout - Adds format validation for --log-format flag The implementation ensures logs are flushed on all exit paths including normal shutdown, errors, and panics via deferred close.
1 parent 7519ee2 commit d5db6cf

File tree

10 files changed

+582
-20
lines changed

10 files changed

+582
-20
lines changed

‎cmd/spicedb/main.go‎

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package main
22

33
import (
44
"errors"
5+
"fmt"
56
"os"
67

78
"github.com/rs/zerolog"
@@ -18,11 +19,20 @@ import (
1819
)
1920

2021
func main() {
22+
os.Exit(run())
23+
}
24+
25+
func run() int {
2126
memoryprotection.InitDefaultMemoryUsageProvider()
2227

2328
// Set up root logger
2429
// This will typically be overwritten by the logging setup for a given command.
2530
log.SetGlobalLogger(zerolog.New(os.Stderr).Level(zerolog.InfoLevel))
31+
defer func() {
32+
if err := log.Close(); err != nil {
33+
fmt.Fprintf(os.Stderr, "warning: failed to flush logs: %v\n", err)
34+
}
35+
}()
2636

2737
// Enable Kubernetes gRPC resolver
2838
kuberesolver.RegisterInCluster()
@@ -33,7 +43,8 @@ func main() {
3343
// Build the complete command structure
3444
rootCmd, err := cmd.BuildRootCommand()
3545
if err != nil {
36-
log.Fatal().Err(err).Msg("failed to build root command")
46+
log.Error().Err(err).Msg("failed to build root command")
47+
return 1
3748
}
3849

3950
if err := rootCmd.Execute(); err != nil {
@@ -42,8 +53,10 @@ func main() {
4253
}
4354
var termErr spiceerrors.TerminationError
4455
if errors.As(err, &termErr) {
45-
os.Exit(termErr.ExitCode())
56+
return termErr.ExitCode()
4657
}
47-
os.Exit(1)
58+
return 1
4859
}
60+
61+
return 0
4962
}

‎internal/logging/logger.go‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@ package logging
22

33
import (
44
"context"
5+
"io"
6+
"sync"
57

68
"github.com/go-logr/zerologr"
79
"github.com/rs/zerolog"
@@ -10,6 +12,59 @@ import (
1012

1113
var Logger zerolog.Logger
1214

15+
type closerHolder struct {
16+
mu sync.Mutex
17+
closer io.Closer
18+
closed bool
19+
}
20+
21+
func (ch *closerHolder) Set(closer io.Closer) {
22+
var toClose io.Closer
23+
24+
ch.mu.Lock()
25+
if ch.closed {
26+
toClose = closer
27+
ch.mu.Unlock()
28+
if toClose != nil {
29+
_ = toClose.Close()
30+
}
31+
return
32+
}
33+
34+
// Don't close if the same closer is being set again
35+
if ch.closer != closer {
36+
toClose = ch.closer
37+
}
38+
ch.closer = closer
39+
ch.mu.Unlock()
40+
41+
if toClose != nil {
42+
_ = toClose.Close()
43+
}
44+
}
45+
46+
func (ch *closerHolder) Close() error {
47+
var toClose io.Closer
48+
49+
ch.mu.Lock()
50+
if ch.closed {
51+
ch.mu.Unlock()
52+
return nil
53+
}
54+
55+
ch.closed = true
56+
toClose = ch.closer
57+
ch.closer = nil
58+
ch.mu.Unlock()
59+
60+
if toClose != nil {
61+
return toClose.Close()
62+
}
63+
return nil
64+
}
65+
66+
var globalCloser closerHolder
67+
1368
func init() {
1469
SetGlobalLogger(zerolog.Nop())
1570
logf.SetLogger(zerologr.New(&Logger))
@@ -20,6 +75,15 @@ func SetGlobalLogger(logger zerolog.Logger) {
2075
zerolog.DefaultContextLogger = &Logger
2176
}
2277

78+
func SetGlobalLoggerWithCloser(logger zerolog.Logger, closer io.Closer) {
79+
SetGlobalLogger(logger)
80+
globalCloser.Set(closer)
81+
}
82+
83+
// Close flushes and releases resources owned by the globally configured logger.
84+
// It is safe to call multiple times.
85+
func Close() error { return globalCloser.Close() }
86+
2387
func With() zerolog.Context { return Logger.With() }
2488

2589
func Err(err error) *zerolog.Event { return Logger.Err(err) }

‎internal/logging/logger_test.go‎

Lines changed: 215 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,215 @@
1+
package logging
2+
3+
import (
4+
"errors"
5+
"sync"
6+
"sync/atomic"
7+
"testing"
8+
9+
"github.com/rs/zerolog"
10+
"github.com/stretchr/testify/require"
11+
)
12+
13+
// mockCloser is a test double that tracks Close() calls
14+
type mockCloser struct {
15+
closeCount atomic.Int32
16+
closeErr error
17+
}
18+
19+
func (m *mockCloser) Close() error {
20+
m.closeCount.Add(1)
21+
return m.closeErr
22+
}
23+
24+
func TestCloserHolder_Set(t *testing.T) {
25+
t.Run("sets closer", func(t *testing.T) {
26+
ch := &closerHolder{}
27+
closer := &mockCloser{}
28+
29+
ch.Set(closer)
30+
31+
require.NoError(t, ch.Close())
32+
require.Equal(t, int32(1), closer.closeCount.Load())
33+
})
34+
35+
t.Run("closes previous closer when setting new one", func(t *testing.T) {
36+
ch := &closerHolder{}
37+
first := &mockCloser{}
38+
second := &mockCloser{}
39+
40+
ch.Set(first)
41+
ch.Set(second)
42+
43+
// First should have been closed when second was set
44+
require.Equal(t, int32(1), first.closeCount.Load())
45+
// Second should not be closed yet
46+
require.Equal(t, int32(0), second.closeCount.Load())
47+
48+
// Now close the holder
49+
require.NoError(t, ch.Close())
50+
require.Equal(t, int32(1), second.closeCount.Load())
51+
})
52+
53+
t.Run("does not double-close same closer", func(t *testing.T) {
54+
ch := &closerHolder{}
55+
closer := &mockCloser{}
56+
57+
ch.Set(closer)
58+
ch.Set(closer) // Same closer again
59+
60+
// Should not have been closed yet (same instance)
61+
require.Equal(t, int32(0), closer.closeCount.Load())
62+
63+
require.NoError(t, ch.Close())
64+
require.Equal(t, int32(1), closer.closeCount.Load())
65+
})
66+
67+
t.Run("closes new closer immediately if holder already closed", func(t *testing.T) {
68+
ch := &closerHolder{}
69+
first := &mockCloser{}
70+
second := &mockCloser{}
71+
72+
ch.Set(first)
73+
require.NoError(t, ch.Close())
74+
75+
// Setting after close should immediately close the new closer
76+
ch.Set(second)
77+
require.Equal(t, int32(1), second.closeCount.Load())
78+
})
79+
80+
t.Run("handles nil closer", func(t *testing.T) {
81+
ch := &closerHolder{}
82+
83+
ch.Set(nil)
84+
require.NoError(t, ch.Close())
85+
})
86+
}
87+
88+
func TestCloserHolder_Close(t *testing.T) {
89+
t.Run("is idempotent", func(t *testing.T) {
90+
ch := &closerHolder{}
91+
closer := &mockCloser{}
92+
93+
ch.Set(closer)
94+
95+
require.NoError(t, ch.Close())
96+
require.NoError(t, ch.Close())
97+
require.NoError(t, ch.Close())
98+
99+
// Should only be closed once
100+
require.Equal(t, int32(1), closer.closeCount.Load())
101+
})
102+
103+
t.Run("returns error from closer", func(t *testing.T) {
104+
ch := &closerHolder{}
105+
expectedErr := errors.New("close failed")
106+
closer := &mockCloser{closeErr: expectedErr}
107+
108+
ch.Set(closer)
109+
110+
err := ch.Close()
111+
require.ErrorIs(t, err, expectedErr)
112+
})
113+
114+
t.Run("returns nil when no closer set", func(t *testing.T) {
115+
ch := &closerHolder{}
116+
require.NoError(t, ch.Close())
117+
})
118+
}
119+
120+
func TestCloserHolder_Concurrent(t *testing.T) {
121+
t.Run("concurrent Set calls are safe", func(t *testing.T) {
122+
ch := &closerHolder{}
123+
var wg sync.WaitGroup
124+
closers := make([]*mockCloser, 100)
125+
126+
for i := range closers {
127+
closers[i] = &mockCloser{}
128+
}
129+
130+
// Concurrently set different closers
131+
for i := 0; i < 100; i++ {
132+
wg.Add(1)
133+
go func(idx int) {
134+
defer wg.Done()
135+
ch.Set(closers[idx])
136+
}(i)
137+
}
138+
139+
wg.Wait()
140+
require.NoError(t, ch.Close())
141+
142+
// Count total closes - should be 100 (99 replaced + 1 final close)
143+
var totalCloses int32
144+
for _, c := range closers {
145+
totalCloses += c.closeCount.Load()
146+
}
147+
require.Equal(t, int32(100), totalCloses)
148+
})
149+
150+
t.Run("concurrent Close calls are safe", func(t *testing.T) {
151+
ch := &closerHolder{}
152+
closer := &mockCloser{}
153+
ch.Set(closer)
154+
155+
var wg sync.WaitGroup
156+
for i := 0; i < 100; i++ {
157+
wg.Add(1)
158+
go func() {
159+
defer wg.Done()
160+
_ = ch.Close()
161+
}()
162+
}
163+
164+
wg.Wait()
165+
166+
// Should only be closed once despite concurrent calls
167+
require.Equal(t, int32(1), closer.closeCount.Load())
168+
})
169+
}
170+
171+
func TestSetGlobalLoggerWithCloser(t *testing.T) {
172+
// Save original state
173+
originalLogger := Logger
174+
175+
t.Cleanup(func() {
176+
SetGlobalLogger(originalLogger)
177+
// Reset global closer
178+
globalCloser = closerHolder{}
179+
})
180+
181+
t.Run("sets logger and closer", func(t *testing.T) {
182+
logger := zerolog.Nop()
183+
closer := &mockCloser{}
184+
185+
SetGlobalLoggerWithCloser(logger, closer)
186+
187+
require.NoError(t, Close())
188+
require.Equal(t, int32(1), closer.closeCount.Load())
189+
})
190+
}
191+
192+
func TestClose(t *testing.T) {
193+
// Save original state
194+
originalLogger := Logger
195+
196+
t.Cleanup(func() {
197+
SetGlobalLogger(originalLogger)
198+
// Reset global closer
199+
globalCloser = closerHolder{}
200+
})
201+
202+
t.Run("closes global closer", func(t *testing.T) {
203+
closer := &mockCloser{}
204+
SetGlobalLoggerWithCloser(zerolog.Nop(), closer)
205+
206+
require.NoError(t, Close())
207+
require.Equal(t, int32(1), closer.closeCount.Load())
208+
})
209+
210+
t.Run("is safe to call without closer", func(t *testing.T) {
211+
// Reset to ensure no closer
212+
globalCloser = closerHolder{}
213+
require.NoError(t, Close())
214+
})
215+
}

‎pkg/cmd/lsp.go‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,11 @@ import (
66

77
"github.com/go-logr/zerologr"
88
"github.com/jzelinskie/cobrautil/v2"
9-
"github.com/jzelinskie/cobrautil/v2/cobrazerolog"
10-
"github.com/rs/zerolog"
119
"github.com/spf13/cobra"
1210

1311
"github.com/authzed/spicedb/internal/logging"
1412
"github.com/authzed/spicedb/internal/lsp"
13+
"github.com/authzed/spicedb/pkg/cmd/server"
1514
"github.com/authzed/spicedb/pkg/cmd/termination"
1615
"github.com/authzed/spicedb/pkg/releases"
1716
)
@@ -41,11 +40,7 @@ func NewLSPCommand(programName string, config *LSPConfig) *cobra.Command {
4140
Short: "serve language server protocol",
4241
PreRunE: cobrautil.CommandStack(
4342
cobrautil.SyncViperDotEnvPreRunE(programName, "spicedb.env", zerologr.New(&logging.Logger)),
44-
cobrazerolog.New(
45-
cobrazerolog.WithTarget(func(logger zerolog.Logger) {
46-
logging.SetGlobalLogger(logger)
47-
}),
48-
).RunE(),
43+
server.ConfigureLoggingRunE(),
4944
releases.CheckAndLogRunE(),
5045
),
5146
RunE: termination.PublishError(func(cmd *cobra.Command, args []string) error {

‎pkg/cmd/migrate.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ func migrateRun(cmd *cobra.Command, args []string) error {
8787
var err error
8888
emulatorHost, err := cmd.Flags().GetString("datastore-spanner-emulator-host")
8989
if err != nil {
90-
log.Ctx(cmd.Context()).Fatal().Err(err).Msg("unable to get spanner emulator host")
90+
return fmt.Errorf("unable to get spanner emulator host: %w", err)
9191
}
9292
migrationDriver, err := spannermigrations.NewSpannerDriver(cmd.Context(), dbURL, credFile, emulatorHost)
9393
if err != nil {
@@ -100,7 +100,7 @@ func migrateRun(cmd *cobra.Command, args []string) error {
100100
var err error
101101
tablePrefix, err := cmd.Flags().GetString("datastore-mysql-table-prefix")
102102
if err != nil {
103-
log.Ctx(cmd.Context()).Fatal().Msg(fmt.Sprintf("unable to get table prefix: %s", err))
103+
return fmt.Errorf("unable to get table prefix: %w", err)
104104
}
105105

106106
var credentialsProvider datastore.CredentialsProvider

0 commit comments

Comments
 (0)