66 "fmt"
77 "net/url"
88 "os"
9+ "sync"
910 "time"
1011
1112 "cdr.dev/slog"
@@ -27,13 +28,14 @@ var (
2728 minAgentAPIV2 = "v2.9"
2829)
2930
30- // Coder establishes a connection to the Coder instance located at
31- // coderURL and authenticates using token. It then establishes a
32- // dRPC connection to the Agent API and begins sending logs.
33- // If the version of Coder does not support the Agent API, it will
34- // fall back to using the PatchLogs endpoint.
35- // The returned function is used to block until all logs are sent.
36- func Coder (ctx context.Context , coderURL * url.URL , token string ) (Func , func (), error ) {
31+ // Coder establishes a connection to the Coder instance located at coderURL and
32+ // authenticates using token. It then establishes a dRPC connection to the Agent
33+ // API and begins sending logs. If the version of Coder does not support the
34+ // Agent API, it will fall back to using the PatchLogs endpoint. The closer is
35+ // used to close the logger and to wait at most logSendGracePeriod for logs to
36+ // be sent. Cancelling the context will close the logs immediately without
37+ // waiting for logs to be sent.
38+ func Coder (ctx context.Context , coderURL * url.URL , token string ) (logger Func , closer func (), err error ) {
3739 // To troubleshoot issues, we need some way of logging.
3840 metaLogger := slog .Make (sloghuman .Sink (os .Stderr ))
3941 defer metaLogger .Sync ()
@@ -44,18 +46,39 @@ func Coder(ctx context.Context, coderURL *url.URL, token string) (Func, func(),
4446 }
4547 if semver .Compare (semver .MajorMinor (bi .Version ), minAgentAPIV2 ) < 0 {
4648 metaLogger .Warn (ctx , "Detected Coder version incompatible with AgentAPI v2, falling back to deprecated API" , slog .F ("coder_version" , bi .Version ))
47- sendLogs , flushLogs : =sendLogsV1 (ctx , client , metaLogger .Named ("send_logs_v1" ))
48- return sendLogs , flushLogs , nil
49+ logger , closer = sendLogsV1 (ctx , client , metaLogger .Named ("send_logs_v1" ))
50+ return logger , closer , nil
4951 }
52+
53+ // Create a new context so we can ensure the connection is torn down.
54+ ctx , cancel := context .WithCancel (ctx )
55+ defer func () {
56+ if err != nil {
57+ cancel ()
58+ }
59+ }()
60+ // Note that ctx passed to initRPC will be inherited by the
61+ // underlying connection, nothing we can do about that here.
5062 dac , err := initRPC (ctx , client , metaLogger .Named ("init_rpc" ))
5163 if err != nil {
5264 // Logged externally
5365 return nil , nil , fmt .Errorf ("init coder rpc client: %w" , err )
5466 }
5567 ls := agentsdk .NewLogSender (metaLogger .Named ("coder_log_sender" ))
5668 metaLogger .Warn (ctx , "Sending logs via AgentAPI v2" , slog .F ("coder_version" , bi .Version ))
57- sendLogs , doneFunc := sendLogsV2 (ctx , dac , ls , metaLogger .Named ("send_logs_v2" ))
58- return sendLogs , doneFunc , nil
69+ logger , loggerCloser := sendLogsV2 (ctx , dac , ls , metaLogger .Named ("send_logs_v2" ))
70+ var closeOnce sync.Once
71+ closer = func () {
72+ loggerCloser ()
73+
74+ closeOnce .Do (func () {
75+ // Typically cancel would be after Close, but we want to be
76+ // sure there's nothing that might block on Close.
77+ cancel ()
78+ _ = dac .DRPCConn ().Close ()
79+ })
80+ }
81+ return logger , closer , nil
5982}
6083
6184type coderLogSender interface {
@@ -74,7 +97,7 @@ func initClient(coderURL *url.URL, token string) *agentsdk.Client {
7497func initRPC (ctx context.Context , client * agentsdk.Client , l slog.Logger ) (proto.DRPCAgentClient20 , error ) {
7598 var c proto.DRPCAgentClient20
7699 var err error
77- retryCtx , retryCancel := context .WithTimeout (context . Background () , rpcConnectTimeout )
100+ retryCtx , retryCancel := context .WithTimeout (ctx , rpcConnectTimeout )
78101 defer retryCancel ()
79102 attempts := 0
80103 for r := retry .New (100 * time .Millisecond , time .Second ); r .Wait (retryCtx ); {
@@ -95,65 +118,67 @@ func initRPC(ctx context.Context, client *agentsdk.Client, l slog.Logger) (proto
95118
96119// sendLogsV1 uses the PatchLogs endpoint to send logs.
97120// This is deprecated, but required for backward compatibility with older versions of Coder.
98- func sendLogsV1 (ctx context.Context , client * agentsdk.Client , l slog.Logger ) (Func , func ()) {
121+ func sendLogsV1 (ctx context.Context , client * agentsdk.Client , l slog.Logger ) (logger Func , closer func ()) {
99122 // nolint: staticcheck // required for backwards compatibility
100- sendLogs , flushLogs := agentsdk .LogsSender (agentsdk .ExternalLogSourceID , client .PatchLogs , slog.Logger {})
123+ sendLog , flushAndClose := agentsdk .LogsSender (agentsdk .ExternalLogSourceID , client .PatchLogs , slog.Logger {})
124+ var mu sync.Mutex
101125 return func (lvl Level , msg string , args ... any ) {
102126 log := agentsdk.Log {
103127 CreatedAt : time .Now (),
104128 Output : fmt .Sprintf (msg , args ... ),
105129 Level : codersdk .LogLevel (lvl ),
106130 }
107- if err := sendLogs (ctx , log ); err != nil {
131+ mu .Lock ()
132+ defer mu .Unlock ()
133+ if err := sendLog (ctx , log ); err != nil {
108134 l .Warn (ctx , "failed to send logs to Coder" , slog .Error (err ))
109135 }
110136 }, func () {
111- if err := flushLogs (ctx ); err != nil {
137+ ctx , cancel := context .WithTimeout (ctx , logSendGracePeriod )
138+ defer cancel ()
139+ if err := flushAndClose (ctx ); err != nil {
112140 l .Warn (ctx , "failed to flush logs" , slog .Error (err ))
113141 }
114142 }
115143}
116144
117145// sendLogsV2 uses the v2 agent API to send logs. Only compatibile with coder versions >= 2.9.
118- func sendLogsV2 (ctx context.Context , dest agentsdk.LogDest , ls coderLogSender , l slog.Logger ) (Func , func ()) {
146+ func sendLogsV2 (ctx context.Context , dest agentsdk.LogDest , ls coderLogSender , l slog.Logger ) (logger Func , closer func ()) {
147+ sendCtx , sendCancel := context .WithCancel (ctx )
119148 done := make (chan struct {})
120149 uid := uuid .New ()
121150 go func () {
122151 defer close (done )
123- if err := ls .SendLoop (ctx , dest ); err != nil {
152+ if err := ls .SendLoop (sendCtx , dest ); err != nil {
124153 if ! errors .Is (err , context .Canceled ) {
125154 l .Warn (ctx , "failed to send logs to Coder" , slog .Error (err ))
126155 }
127156 }
128-
129- // Wait for up to 10 seconds for logs to finish sending.
130- sendCtx , sendCancel := context .WithTimeout (context .Background (), logSendGracePeriod )
131- defer sendCancel ()
132- // Try once more to send any pending logs
133- if err := ls .SendLoop (sendCtx , dest ); err != nil {
134- if ! errors .Is (err , context .DeadlineExceeded ) {
135- l .Warn (ctx , "failed to send remaining logs to Coder" , slog .Error (err ))
136- }
137- }
138- ls .Flush (uid )
139- if err := ls .WaitUntilEmpty (sendCtx ); err != nil {
140- if ! errors .Is (err , context .DeadlineExceeded ) {
141- l .Warn (ctx , "log sender did not empty" , slog .Error (err ))
142- }
143- }
144157 }()
145158
146- logFunc := func (l Level , msg string , args ... any ) {
147- ls .Enqueue (uid , agentsdk.Log {
148- CreatedAt : time .Now (),
149- Output : fmt .Sprintf (msg , args ... ),
150- Level : codersdk .LogLevel (l ),
151- })
152- }
159+ var closeOnce sync.Once
160+ return func (l Level , msg string , args ... any ) {
161+ ls .Enqueue (uid , agentsdk.Log {
162+ CreatedAt : time .Now (),
163+ Output : fmt .Sprintf (msg , args ... ),
164+ Level : codersdk .LogLevel (l ),
165+ })
166+ }, func () {
167+ closeOnce .Do (func () {
168+ // Trigger a flush and wait for logs to be sent.
169+ ls .Flush (uid )
170+ ctx , cancel := context .WithTimeout (ctx , logSendGracePeriod )
171+ defer cancel ()
172+ err := ls .WaitUntilEmpty (ctx )
173+ if err != nil {
174+ l .Warn (ctx , "log sender did not empty" , slog .Error (err ))
175+ }
153176
154- doneFunc := func () {
155- <- done
156- }
177+ // Stop the send loop.
178+ sendCancel ()
179+ })
157180
158- return logFunc , doneFunc
181+ // Wait for the send loop to finish.
182+ <- done
183+ }
159184}
0 commit comments