|
| 1 | +package dbtestutil |
| 2 | + |
| 3 | +import ( |
| 4 | +"context" |
| 5 | +"database/sql" |
| 6 | +"fmt" |
| 7 | +"io" |
| 8 | +"os" |
| 9 | +"os/exec" |
| 10 | +"os/signal" |
| 11 | +"time" |
| 12 | + |
| 13 | +"github.com/google/uuid" |
| 14 | +"golang.org/x/xerrors" |
| 15 | + |
| 16 | +"cdr.dev/slog" |
| 17 | +"cdr.dev/slog/sloggers/sloghuman" |
| 18 | +"github.com/coder/retry" |
| 19 | +) |
| 20 | + |
| 21 | +const ( |
| 22 | +cleanerRespOK="OK" |
| 23 | +envCleanerParentUUID="DB_CLEANER_PARENT_UUID" |
| 24 | +) |
| 25 | + |
| 26 | +// startCleaner starts the cleaner in a subprocess. holdThis is an opaque reference that needs to be kept from being |
| 27 | +// garbage collected until we are done with all test databases (e.g. the end of the processs). |
| 28 | +funcstartCleaner(ctx context.Context,parentUUID uuid.UUID) (holdThisany,errerror) { |
| 29 | +cmd:=exec.Command("go","run","github.com/coder/coder/v2/cmd/dbtestcleaner") |
| 30 | +cmd.Env=append(os.Environ(),fmt.Sprintf("%s=%s",envCleanerParentUUID,parentUUID.String())) |
| 31 | + |
| 32 | +// Here we don't actually use the reference to the stdin pipe, because we never write anything to it. When this |
| 33 | +// process exits, the pipe is closed by the OS and this triggers the cleaner to do its cleaning work. But, we do |
| 34 | +// need to hang on to a reference to it so that it doesn't get garbage collected and trigger cleanup early. |
| 35 | +stdin,err:=cmd.StdinPipe() |
| 36 | +iferr!=nil { |
| 37 | +returnnil,xerrors.Errorf("failed to open stdin pipe: %w",err) |
| 38 | +} |
| 39 | +stdout,err:=cmd.StdoutPipe() |
| 40 | +iferr!=nil { |
| 41 | +returnnil,xerrors.Errorf("failed to open stdout pipe: %w",err) |
| 42 | +} |
| 43 | +// uncomment this to see log output from the cleaner |
| 44 | +//cmd.Stderr = os.Stderr |
| 45 | +err=cmd.Start() |
| 46 | +iferr!=nil { |
| 47 | +returnnil,xerrors.Errorf("failed to start broker: %w",err) |
| 48 | +} |
| 49 | +outCh:=make(chan []byte,1) |
| 50 | +errCh:=make(chanerror,1) |
| 51 | +gofunc() { |
| 52 | +buf:=make([]byte,1024) |
| 53 | +n,readErr:=stdout.Read(buf) |
| 54 | +ifreadErr!=nil { |
| 55 | +errCh<-readErr |
| 56 | +return |
| 57 | +} |
| 58 | +outCh<-buf[:n] |
| 59 | +}() |
| 60 | +select { |
| 61 | +case<-ctx.Done(): |
| 62 | +_=cmd.Process.Kill() |
| 63 | +returnnil,ctx.Err() |
| 64 | +caseerr:=<-errCh: |
| 65 | +returnnil,xerrors.Errorf("failed to read db test cleaner output: %w",err) |
| 66 | +caseout:=<-outCh: |
| 67 | +ifstring(out)!=cleanerRespOK { |
| 68 | +returnnil,xerrors.Errorf("db test cleaner error: %s",string(out)) |
| 69 | +} |
| 70 | +returnstdin,nil |
| 71 | +} |
| 72 | +} |
| 73 | + |
| 74 | +typecleanerstruct { |
| 75 | +parentUUID uuid.UUID |
| 76 | +logger slog.Logger |
| 77 | +db*sql.DB |
| 78 | +} |
| 79 | + |
| 80 | +func (c*cleaner)init(ctx context.Context)error { |
| 81 | +varerrerror |
| 82 | +parentUUIDStr:=os.Getenv(envCleanerParentUUID) |
| 83 | +c.parentUUID,err=uuid.Parse(parentUUIDStr) |
| 84 | +iferr!=nil { |
| 85 | +returnxerrors.Errorf("failed to parse parent UUID '%s': %w",parentUUIDStr,err) |
| 86 | +} |
| 87 | +c.logger=slog.Make(sloghuman.Sink(os.Stderr)). |
| 88 | +Named("dbtestcleaner"). |
| 89 | +Leveled(slog.LevelDebug). |
| 90 | +With(slog.F("parent_uuid",parentUUIDStr)) |
| 91 | + |
| 92 | +// TODO: support configurable username, password & port if required |
| 93 | +dsn:=fmt.Sprintf("postgres://postgres:postgres@localhost:5432/%s?sslmode=disable",CoderTestingDBName) |
| 94 | +c.db,err=sql.Open("postgres",dsn) |
| 95 | +iferr!=nil { |
| 96 | +returnxerrors.Errorf("couldn't open DB: %w",err) |
| 97 | +} |
| 98 | +forr:=retry.New(10*time.Millisecond,500*time.Millisecond);r.Wait(ctx); { |
| 99 | +err=c.db.PingContext(ctx) |
| 100 | +iferr==nil { |
| 101 | +returnnil |
| 102 | +} |
| 103 | +c.logger.Error(ctx,"failed to ping DB",slog.Error(err)) |
| 104 | +} |
| 105 | +returnctx.Err() |
| 106 | +} |
| 107 | + |
| 108 | +// waitAndClean waits for stdin to close then attempts to clean up any test databases with our parent's UUID. This |
| 109 | +// is best-effort. If we hit an error we exit. |
| 110 | +// |
| 111 | +// We log to stderr for debugging, but we don't expect this output to normally be available since the parent has |
| 112 | +// exited. Uncomment the line `cmd.Stderr = os.Stderr` in startCleaner() to see this output. |
| 113 | +func (c*cleaner)waitAndClean() { |
| 114 | +c.logger.Debug(context.Background(),"waiting for stdin to close") |
| 115 | +_,_=io.ReadAll(os.Stdin)// here we're just waiting for stdin to close |
| 116 | +c.logger.Debug(context.Background(),"stdin closed") |
| 117 | +rows,err:=c.db.Query( |
| 118 | +"SELECT name FROM test_databases WHERE process_uuid = $1 AND dropped_at IS NULL", |
| 119 | +c.parentUUID, |
| 120 | +) |
| 121 | +iferr!=nil { |
| 122 | +c.logger.Error(context.Background(),"error querying",slog.Error(err)) |
| 123 | +return |
| 124 | +} |
| 125 | +deferrows.Close() |
| 126 | +names:=make([]string,0) |
| 127 | +forrows.Next() { |
| 128 | +varnamestring |
| 129 | +iferr:=rows.Scan(&name);err!=nil { |
| 130 | +continue |
| 131 | +} |
| 132 | +names=append(names,name) |
| 133 | +} |
| 134 | +rows.Close() |
| 135 | +c.logger.Debug(context.Background(),"queried names",slog.F("names",names)) |
| 136 | +for_,name:=rangenames { |
| 137 | +_,err:=c.db.Exec(fmt.Sprintf("DROP DATABASE IF EXISTS %s",name)) |
| 138 | +iferr!=nil { |
| 139 | +c.logger.Error(context.Background(),"error dropping database",slog.Error(err),slog.F("name",name)) |
| 140 | +return |
| 141 | +} |
| 142 | +_,err=c.db.Exec("UPDATE test_databases SET dropped_at = CURRENT_TIMESTAMP WHERE name = $1",name) |
| 143 | +iferr!=nil { |
| 144 | +c.logger.Error(context.Background(),"error dropping database",slog.Error(err),slog.F("name",name)) |
| 145 | +return |
| 146 | +} |
| 147 | +} |
| 148 | +c.logger.Debug(context.Background(),"finished cleaning") |
| 149 | +} |
| 150 | + |
| 151 | +// RunCleaner runs the test database cleaning process. It takes no arguments but uses stdio and environment variables |
| 152 | +// for its operation. It is designed to be launched as the only task of a `main()` process, but is included in this |
| 153 | +// package to share constants with the parent code that launches it above. |
| 154 | +// |
| 155 | +// The cleaner is designed to run in a separate process from the main test suite, connected over stdio. If the main test |
| 156 | +// process ends (panics, times out, or is killed) without explicitly discarding the databases it clones, the cleaner |
| 157 | +// removes them so they don't leak beyond the test session. c.f. https://github.com/coder/internal/issues/927 |
| 158 | +funcRunCleaner() { |
| 159 | +c:=cleaner{} |
| 160 | +ctx,cancel:=context.WithTimeout(context.Background(),10*time.Second) |
| 161 | +defercancel() |
| 162 | +// canceling a test via the IDE sends us an interrupt signal. We only want to process that signal during init. After |
| 163 | +// we want to ignore the signal and do our cleaning. |
| 164 | +signalCtx,signalCancel:=signal.NotifyContext(ctx,os.Interrupt) |
| 165 | +defersignalCancel() |
| 166 | +err:=c.init(signalCtx) |
| 167 | +iferr!=nil { |
| 168 | +_,_=fmt.Fprintf(os.Stdout,"failed to init: %s",err.Error()) |
| 169 | +_=os.Stdout.Close() |
| 170 | +return |
| 171 | +} |
| 172 | +_,_=fmt.Fprint(os.Stdout,cleanerRespOK) |
| 173 | +_=os.Stdout.Close() |
| 174 | +c.waitAndClean() |
| 175 | +} |