Skip to content

Commit c077c72

Browse files
EmyrkEmyrk
committed
Implement workspace proxy cmd
1 parent fd7dd30 commit c077c72

File tree

2 files changed

+178
-57
lines changed

2 files changed

+178
-57
lines changed

‎cli/server.go

Lines changed: 73 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -33,14 +33,15 @@ import (
3333
"sync/atomic"
3434
"time"
3535

36+
"github.com/prometheus/client_golang/prometheus/collectors"
37+
"github.com/prometheus/client_golang/prometheus/promhttp"
38+
3639
"github.com/coreos/go-oidc/v3/oidc"
3740
"github.com/coreos/go-systemd/daemon"
3841
embeddedpostgres "github.com/fergusstrange/embedded-postgres"
3942
"github.com/google/go-github/v43/github"
4043
"github.com/google/uuid"
4144
"github.com/prometheus/client_golang/prometheus"
42-
"github.com/prometheus/client_golang/prometheus/collectors"
43-
"github.com/prometheus/client_golang/prometheus/promhttp"
4445
"github.com/spf13/afero"
4546
"go.opentelemetry.io/otel/trace"
4647
"golang.org/x/mod/semver"
@@ -660,33 +661,33 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
660661
defer options.Telemetry.Close()
661662
}
662663

663-
// This prevents the pprof import from being accidentally deleted.
664-
_ = pprof.Handler
665-
if cfg.Pprof.Enable {
666-
//nolint:revive
667-
defer serveHandler(ctx, logger, nil, cfg.Pprof.Address.String(), "pprof")()
668-
}
669-
if cfg.Prometheus.Enable {
670-
options.PrometheusRegistry.MustRegister(collectors.NewGoCollector())
671-
options.PrometheusRegistry.MustRegister(collectors.NewProcessCollector(collectors.ProcessCollectorOpts{}))
672-
673-
closeUsersFunc, err := prometheusmetrics.ActiveUsers(ctx, options.PrometheusRegistry, options.Database, 0)
674-
if err != nil {
675-
return xerrors.Errorf("register active users prometheus metric: %w", err)
676-
}
677-
defer closeUsersFunc()
678-
679-
closeWorkspacesFunc, err := prometheusmetrics.Workspaces(ctx, options.PrometheusRegistry, options.Database, 0)
680-
if err != nil {
681-
return xerrors.Errorf("register workspaces prometheus metric: %w", err)
682-
}
683-
defer closeWorkspacesFunc()
684-
685-
//nolint:revive
686-
defer serveHandler(ctx, logger, promhttp.InstrumentMetricHandler(
687-
options.PrometheusRegistry, promhttp.HandlerFor(options.PrometheusRegistry, promhttp.HandlerOpts{}),
688-
), cfg.Prometheus.Address.String(), "prometheus")()
689-
}
664+
//// This prevents the pprof import from being accidentally deleted.
665+
//_ = pprof.Handler
666+
//if cfg.Pprof.Enable {
667+
// //nolint:revive
668+
// defer serveHandler(ctx, logger, nil, cfg.Pprof.Address.String(), "pprof")()
669+
//}
670+
//if cfg.Prometheus.Enable {
671+
// options.PrometheusRegistry.MustRegister(collectors.NewGoCollector())
672+
// options.PrometheusRegistry.MustRegister(collectors.NewProcessCollector(collectors.ProcessCollectorOpts{}))
673+
//
674+
// closeUsersFunc, err := prometheusmetrics.ActiveUsers(ctx, options.PrometheusRegistry, options.Database, 0)
675+
// if err != nil {
676+
// return xerrors.Errorf("register active users prometheus metric: %w", err)
677+
// }
678+
// defer closeUsersFunc()
679+
//
680+
// closeWorkspacesFunc, err := prometheusmetrics.Workspaces(ctx, options.PrometheusRegistry, options.Database, 0)
681+
// if err != nil {
682+
// return xerrors.Errorf("register workspaces prometheus metric: %w", err)
683+
// }
684+
// defer closeWorkspacesFunc()
685+
//
686+
// //nolint:revive
687+
// defer serveHandler(ctx, logger, promhttp.InstrumentMetricHandler(
688+
// options.PrometheusRegistry, promhttp.HandlerFor(options.PrometheusRegistry, promhttp.HandlerOpts{}),
689+
// ), cfg.Prometheus.Address.String(), "prometheus")()
690+
//}
690691

691692
if cfg.Swagger.Enable {
692693
options.SwaggerEndpoint = cfg.Swagger.Enable.Value()
@@ -701,6 +702,18 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
701702
}
702703

703704
if cfg.Prometheus.Enable {
705+
closeUsersFunc, err := prometheusmetrics.ActiveUsers(ctx, options.PrometheusRegistry, options.Database, 0)
706+
if err != nil {
707+
return xerrors.Errorf("register active users prometheus metric: %w", err)
708+
}
709+
defer closeUsersFunc()
710+
711+
closeWorkspacesFunc, err := prometheusmetrics.Workspaces(ctx, options.PrometheusRegistry, options.Database, 0)
712+
if err != nil {
713+
return xerrors.Errorf("register workspaces prometheus metric: %w", err)
714+
}
715+
defer closeWorkspacesFunc()
716+
704717
// Agent metrics require reference to the tailnet coordinator, so must be initiated after Coder API.
705718
closeAgentsFunc, err := prometheusmetrics.Agents(ctx, logger, options.PrometheusRegistry, coderAPI.Database, &coderAPI.TailnetCoordinator, options.DERPMap, coderAPI.Options.AgentInactiveDisconnectTimeout, 0)
706719
if err != nil {
@@ -720,7 +733,6 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
720733
},
721734
}
722735
}
723-
defer client.HTTPClient.CloseIdleConnections()
724736

725737
// This is helpful for tests, but can be silently ignored.
726738
// Coder may be ran as users that don't have permission to write in the homedir,
@@ -819,6 +831,11 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
819831
}
820832
}()
821833

834+
autobuildPoller := time.NewTicker(cfg.AutobuildPollInterval.Value())
835+
defer autobuildPoller.Stop()
836+
autobuildExecutor := executor.New(ctx, options.Database, coderAPI.TemplateScheduleStore, logger, autobuildPoller.C)
837+
autobuildExecutor.Run()
838+
822839
cliui.Infof(inv.Stdout, "\n==> Logs will stream in below (press ctrl+c to gracefully exit):")
823840

824841
// Updates the systemd status from activating to activated.
@@ -827,11 +844,6 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
827844
return xerrors.Errorf("notify systemd: %w", err)
828845
}
829846

830-
autobuildPoller := time.NewTicker(cfg.AutobuildPollInterval.Value())
831-
defer autobuildPoller.Stop()
832-
autobuildExecutor := executor.New(ctx, options.Database, coderAPI.TemplateScheduleStore, logger, autobuildPoller.C)
833-
autobuildExecutor.Run()
834-
835847
// Currently there is no way to ask the server to shut
836848
// itself down, so any exit signal will result in a non-zero
837849
// exit of the server.
@@ -1214,9 +1226,10 @@ type CommonServerCmd struct {
12141226
HTTPServers *HTTPServers
12151227
HTTPClient *http.Client
12161228
LocalURL *url.URL
1217-
Logger slog.Logger
12181229

1219-
Tracer trace.TracerProvider
1230+
Logger slog.Logger
1231+
PrometheusRegistry *prometheus.Registry
1232+
Tracer trace.TracerProvider
12201233
// SQLDriver is the driver that the tracer is active on.
12211234
SQLDriver string
12221235

@@ -1328,12 +1341,6 @@ func SetupServerCmd(inv *clibase.Invocation, cfg *codersdk.DeploymentValues) (_
13281341
// A newline is added before for visibility in terminal output.
13291342
cliui.Infof(inv.Stdout, "\nView the Web UI: %s", cfg.AccessURL.String())
13301343

1331-
// Used for zero-trust instance identity with Google Cloud.
1332-
googleTokenValidator, err := idtoken.NewValidator(ctx, option.WithoutAuthentication())
1333-
if err != nil {
1334-
return nil, err
1335-
}
1336-
13371344
c.AppHostname = cfg.WildcardAccessURL.String()
13381345
if c.AppHostname != "" {
13391346
c.AppHostnameRegex, err = httpapi.CompileHostnamePattern(c.AppHostname)
@@ -1347,6 +1354,29 @@ func SetupServerCmd(inv *clibase.Invocation, cfg *codersdk.DeploymentValues) (_
13471354
return nil, xerrors.Errorf("parse real ip config: %w", err)
13481355
}
13491356

1357+
if cfg.Pprof.Enable {
1358+
// This prevents the pprof import from being accidentally deleted.
1359+
// pprof has an init function that attaches itself to the default handler.
1360+
// By passing a nil handler to 'serverHandler', it will automatically use
1361+
// the default, which has pprof attached.
1362+
_ = pprof.Handler
1363+
//nolint:revive
1364+
closeFunc := serveHandler(ctx, logger, nil, cfg.Pprof.Address.String(), "pprof")
1365+
c.addClose(closeFunc)
1366+
}
1367+
1368+
c.PrometheusRegistry = prometheus.NewRegistry()
1369+
if cfg.Prometheus.Enable {
1370+
c.PrometheusRegistry.MustRegister(collectors.NewGoCollector())
1371+
c.PrometheusRegistry.MustRegister(collectors.NewProcessCollector(collectors.ProcessCollectorOpts{}))
1372+
1373+
//nolint:revive
1374+
closeFunc := serveHandler(ctx, logger, promhttp.InstrumentMetricHandler(
1375+
c.PrometheusRegistry, promhttp.HandlerFor(c.PrometheusRegistry, promhttp.HandlerOpts{}),
1376+
), cfg.Prometheus.Address.String(), "prometheus")
1377+
c.addClose(closeFunc)
1378+
}
1379+
13501380
return c, nil
13511381
}
13521382

‎enterprise/cli/workspaceproxy.go

Lines changed: 105 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,19 @@
11
package cli
22

33
import (
4+
"context"
5+
"fmt"
46
"io"
57
"log"
8+
"net"
69
"net/http"
10+
"runtime/pprof"
11+
"time"
12+
13+
"github.com/coreos/go-systemd/daemon"
14+
15+
"github.com/coder/coder/cli/cliui"
16+
"golang.org/x/xerrors"
717

818
"github.com/coder/coder/cli"
919
"github.com/coder/coder/coderd/workspaceapps"
@@ -54,23 +64,30 @@ func (r *RootCmd) proxyServer() *clibase.Cmd {
5464
if err != nil {
5565
return err
5666
}
67+
defer scd.Close()
68+
ctx := scd.Ctx
5769

5870
proxy, err := wsproxy.New(&wsproxy.Options{
59-
Logger: scd.Logger,
71+
Logger: scd.Logger,
72+
// TODO: PrimaryAccessURL
6073
PrimaryAccessURL: nil,
61-
AccessURL: nil,
74+
AccessURL: cfg.AccessURL.Value(),
6275
AppHostname: scd.AppHostname,
6376
AppHostnameRegex: scd.AppHostnameRegex,
6477
RealIPConfig: scd.RealIPConfig,
6578
AppSecurityKey: workspaceapps.SecurityKey{},
6679
Tracing: scd.Tracer,
67-
PrometheusRegistry: nil,
68-
APIRateLimit: 0,
69-
SecureAuthCookie: false,
70-
DisablePathApps: false,
71-
ProxySessionToken: "",
80+
PrometheusRegistry: scd.PrometheusRegistry,
81+
APIRateLimit: int(cfg.RateLimit.API.Value()),
82+
SecureAuthCookie: cfg.SecureAuthCookie.Value(),
83+
// TODO: DisablePathApps
84+
DisablePathApps: false,
85+
// TODO: ProxySessionToken
86+
ProxySessionToken: "",
7287
})
7388

89+
shutdownConnsCtx, shutdownConns := context.WithCancel(ctx)
90+
defer shutdownConns()
7491
// ReadHeaderTimeout is purposefully not enabled. It caused some
7592
// issues with websockets over the dev tunnel.
7693
// See: https://github.com/coder/coder/pull/3730
@@ -81,16 +98,90 @@ func (r *RootCmd) proxyServer() *clibase.Cmd {
8198
// https://github.com/hashicorp/vault/blob/e2490059d0711635e529a4efcbaa1b26998d6e1c/command/server.go#L2714
8299
ErrorLog: log.New(io.Discard, "", 0),
83100
Handler: proxy.Handler,
84-
//BaseContext: func(_ net.Listener) context.Context {
85-
// return shutdownConnsCtx
86-
//},
101+
BaseContext: func(_ net.Listener) context.Context {
102+
return shutdownConnsCtx
103+
},
87104
}
88-
//defer func() {
89-
// _ = shutdownWithTimeout(httpServer.Shutdown, 5*time.Second)
90-
//}()
105+
defer func() {
106+
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
107+
defer cancel()
108+
_ = httpServer.Shutdown(ctx)
109+
}()
91110

92111
// TODO: So this obviously is not going to work well.
93-
return scd.HTTPServers.Serve(httpServer)
112+
errCh := make(chan error, 1)
113+
pprof.Do(ctx, pprof.Labels("service", "workspace-proxy"), func(ctx context.Context) {
114+
errCh <- scd.HTTPServers.Serve(httpServer)
115+
})
116+
117+
cliui.Infof(inv.Stdout, "\n==> Logs will stream in below (press ctrl+c to gracefully exit):")
118+
119+
// Updates the systemd status from activating to activated.
120+
_, err = daemon.SdNotify(false, daemon.SdNotifyReady)
121+
if err != nil {
122+
return xerrors.Errorf("notify systemd: %w", err)
123+
}
124+
125+
// Currently there is no way to ask the server to shut
126+
// itself down, so any exit signal will result in a non-zero
127+
// exit of the server.
128+
var exitErr error
129+
select {
130+
case exitErr = <-errCh:
131+
fmt.Println("As")
132+
case <-scd.NotifyCtx.Done():
133+
exitErr = scd.NotifyCtx.Err()
134+
_, _ = fmt.Fprintln(inv.Stdout, cliui.Styles.Bold.Render(
135+
"Interrupt caught, gracefully exiting. Use ctrl+\\ to force quit",
136+
))
137+
}
138+
139+
if exitErr != nil && !xerrors.Is(exitErr, context.Canceled) {
140+
cliui.Errorf(inv.Stderr, "Unexpected error, shutting down server: %s\n", exitErr)
141+
}
142+
143+
// Begin clean shut down stage, we try to shut down services
144+
// gracefully in an order that gives the best experience.
145+
// This procedure should not differ greatly from the order
146+
// of `defer`s in this function, but allows us to inform
147+
// the user about what's going on and handle errors more
148+
// explicitly.
149+
150+
_, err = daemon.SdNotify(false, daemon.SdNotifyStopping)
151+
if err != nil {
152+
cliui.Errorf(inv.Stderr, "Notify systemd failed: %s", err)
153+
}
154+
155+
// Stop accepting new connections without interrupting
156+
// in-flight requests, give in-flight requests 5 seconds to
157+
// complete.
158+
cliui.Info(inv.Stdout, "Shutting down API server..."+"\n")
159+
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
160+
defer cancel()
161+
err = httpServer.Shutdown(ctx)
162+
if err != nil {
163+
cliui.Errorf(inv.Stderr, "API server shutdown took longer than 3s: %s\n", err)
164+
} else {
165+
cliui.Info(inv.Stdout, "Gracefully shut down API server\n")
166+
}
167+
// Cancel any remaining in-flight requests.
168+
shutdownConns()
169+
170+
// Trigger context cancellation for any remaining services.
171+
scd.Close()
172+
173+
switch {
174+
case xerrors.Is(exitErr, context.DeadlineExceeded):
175+
cliui.Warnf(inv.Stderr, "Graceful shutdown timed out")
176+
// Errors here cause a significant number of benign CI failures.
177+
return nil
178+
case xerrors.Is(exitErr, context.Canceled):
179+
return nil
180+
case exitErr != nil:
181+
return xerrors.Errorf("graceful shutdown: %w", exitErr)
182+
default:
183+
return nil
184+
}
94185
},
95186
}
96187

0 commit comments

Comments
 (0)