5
5
"database/sql"
6
6
"encoding/json"
7
7
"fmt"
8
+ "io"
8
9
"net/http"
9
10
"net/http/cookiejar"
10
11
"net/url"
@@ -104,18 +105,18 @@ func (c *Client) provisionerJobLogsBefore(ctx context.Context, path string, befo
104
105
}
105
106
106
107
// provisionerJobLogsAfter streams logs that occurred after a specific time.
107
- func (c * Client ) provisionerJobLogsAfter (ctx context.Context , path string , after time.Time ) (<- chan ProvisionerJobLog , error ) {
108
+ func (c * Client ) provisionerJobLogsAfter (ctx context.Context , path string , after time.Time ) (<- chan ProvisionerJobLog , io. Closer , error ) {
108
109
afterQuery := ""
109
110
if ! after .IsZero () {
110
111
afterQuery = fmt .Sprintf ("&after=%d" , after .UTC ().UnixMilli ())
111
112
}
112
113
followURL , err := c .URL .Parse (fmt .Sprintf ("%s?follow%s" , path , afterQuery ))
113
114
if err != nil {
114
- return nil , err
115
+ return nil , nil , err
115
116
}
116
117
jar , err := cookiejar .New (nil )
117
118
if err != nil {
118
- return nil , xerrors .Errorf ("create cookie jar: %w" , err )
119
+ return nil , nil , xerrors .Errorf ("create cookie jar: %w" , err )
119
120
}
120
121
jar .SetCookies (followURL , []* http.Cookie {{
121
122
Name : SessionTokenKey ,
@@ -129,11 +130,13 @@ func (c *Client) provisionerJobLogsAfter(ctx context.Context, path string, after
129
130
CompressionMode : websocket .CompressionDisabled ,
130
131
})
131
132
if err != nil {
132
- return nil , readBodyAsError (res )
133
+ return nil , nil , readBodyAsError (res )
133
134
}
134
135
logs := make (chan ProvisionerJobLog )
135
136
decoder := json .NewDecoder (websocket .NetConn (ctx , conn , websocket .MessageText ))
137
+ closed := make (chan struct {})
136
138
go func () {
139
+ defer close (closed )
137
140
defer close (logs )
138
141
defer conn .Close (websocket .StatusGoingAway , "" )
139
142
var log ProvisionerJobLog
@@ -149,5 +152,9 @@ func (c *Client) provisionerJobLogsAfter(ctx context.Context, path string, after
149
152
}
150
153
}
151
154
}()
152
- return logs , nil
155
+ return logs , closeFunc (func () error {
156
+ _ = conn .Close (websocket .StatusNormalClosure , "" )
157
+ <- closed
158
+ return nil
159
+ }), nil
153
160
}
0 commit comments