github.com/pdmccormick/importable-docker-buildx@v0.0.0-20240426161518-e47091289030/controller/remote/server.go (about) 1 package remote 2 3 import ( 4 "context" 5 "io" 6 "sync" 7 "sync/atomic" 8 "time" 9 10 "github.com/docker/buildx/build" 11 controllererrors "github.com/docker/buildx/controller/errdefs" 12 "github.com/docker/buildx/controller/pb" 13 "github.com/docker/buildx/controller/processes" 14 "github.com/docker/buildx/util/ioset" 15 "github.com/docker/buildx/util/progress" 16 "github.com/docker/buildx/version" 17 "github.com/moby/buildkit/client" 18 "github.com/pkg/errors" 19 "golang.org/x/sync/errgroup" 20 ) 21 22 type BuildFunc func(ctx context.Context, options *pb.BuildOptions, stdin io.Reader, progress progress.Writer) (resp *client.SolveResponse, res *build.ResultHandle, err error) 23 24 func NewServer(buildFunc BuildFunc) *Server { 25 return &Server{ 26 buildFunc: buildFunc, 27 } 28 } 29 30 type Server struct { 31 buildFunc BuildFunc 32 session map[string]*session 33 sessionMu sync.Mutex 34 } 35 36 type session struct { 37 buildOnGoing atomic.Bool 38 statusChan chan *pb.StatusResponse 39 cancelBuild func() 40 buildOptions *pb.BuildOptions 41 inputPipe *io.PipeWriter 42 43 result *build.ResultHandle 44 45 processes *processes.Manager 46 } 47 48 func (s *session) cancelRunningProcesses() { 49 s.processes.CancelRunningProcesses() 50 } 51 52 func (m *Server) ListProcesses(ctx context.Context, req *pb.ListProcessesRequest) (res *pb.ListProcessesResponse, err error) { 53 m.sessionMu.Lock() 54 defer m.sessionMu.Unlock() 55 s, ok := m.session[req.Ref] 56 if !ok { 57 return nil, errors.Errorf("unknown ref %q", req.Ref) 58 } 59 res = new(pb.ListProcessesResponse) 60 res.Infos = append(res.Infos, s.processes.ListProcesses()...) 61 return res, nil 62 } 63 64 func (m *Server) DisconnectProcess(ctx context.Context, req *pb.DisconnectProcessRequest) (res *pb.DisconnectProcessResponse, err error) { 65 m.sessionMu.Lock() 66 defer m.sessionMu.Unlock() 67 s, ok := m.session[req.Ref] 68 if !ok { 69 return nil, errors.Errorf("unknown ref %q", req.Ref) 70 } 71 return res, s.processes.DeleteProcess(req.ProcessID) 72 } 73 74 func (m *Server) Info(ctx context.Context, req *pb.InfoRequest) (res *pb.InfoResponse, err error) { 75 return &pb.InfoResponse{ 76 BuildxVersion: &pb.BuildxVersion{ 77 Package: version.Package, 78 Version: version.Version, 79 Revision: version.Revision, 80 }, 81 }, nil 82 } 83 84 func (m *Server) List(ctx context.Context, req *pb.ListRequest) (res *pb.ListResponse, err error) { 85 keys := make(map[string]struct{}) 86 87 m.sessionMu.Lock() 88 for k := range m.session { 89 keys[k] = struct{}{} 90 } 91 m.sessionMu.Unlock() 92 93 var keysL []string 94 for k := range keys { 95 keysL = append(keysL, k) 96 } 97 return &pb.ListResponse{ 98 Keys: keysL, 99 }, nil 100 } 101 102 func (m *Server) Disconnect(ctx context.Context, req *pb.DisconnectRequest) (res *pb.DisconnectResponse, err error) { 103 key := req.Ref 104 if key == "" { 105 return nil, errors.New("disconnect: empty key") 106 } 107 108 m.sessionMu.Lock() 109 if s, ok := m.session[key]; ok { 110 if s.cancelBuild != nil { 111 s.cancelBuild() 112 } 113 s.cancelRunningProcesses() 114 if s.result != nil { 115 s.result.Done() 116 } 117 } 118 delete(m.session, key) 119 m.sessionMu.Unlock() 120 121 return &pb.DisconnectResponse{}, nil 122 } 123 124 func (m *Server) Close() error { 125 m.sessionMu.Lock() 126 for k := range m.session { 127 if s, ok := m.session[k]; ok { 128 if s.cancelBuild != nil { 129 s.cancelBuild() 130 } 131 s.cancelRunningProcesses() 132 } 133 } 134 m.sessionMu.Unlock() 135 return nil 136 } 137 138 func (m *Server) Inspect(ctx context.Context, req *pb.InspectRequest) (*pb.InspectResponse, error) { 139 ref := req.Ref 140 if ref == "" { 141 return nil, errors.New("inspect: empty key") 142 } 143 var bo *pb.BuildOptions 144 m.sessionMu.Lock() 145 if s, ok := m.session[ref]; ok { 146 bo = s.buildOptions 147 } else { 148 m.sessionMu.Unlock() 149 return nil, errors.Errorf("inspect: unknown key %v", ref) 150 } 151 m.sessionMu.Unlock() 152 return &pb.InspectResponse{Options: bo}, nil 153 } 154 155 func (m *Server) Build(ctx context.Context, req *pb.BuildRequest) (*pb.BuildResponse, error) { 156 ref := req.Ref 157 if ref == "" { 158 return nil, errors.New("build: empty key") 159 } 160 161 // Prepare status channel and session 162 m.sessionMu.Lock() 163 if m.session == nil { 164 m.session = make(map[string]*session) 165 } 166 s, ok := m.session[ref] 167 if ok { 168 if !s.buildOnGoing.CompareAndSwap(false, true) { 169 m.sessionMu.Unlock() 170 return &pb.BuildResponse{}, errors.New("build ongoing") 171 } 172 s.cancelRunningProcesses() 173 s.result = nil 174 } else { 175 s = &session{} 176 s.buildOnGoing.Store(true) 177 } 178 179 s.processes = processes.NewManager() 180 statusChan := make(chan *pb.StatusResponse) 181 s.statusChan = statusChan 182 inR, inW := io.Pipe() 183 defer inR.Close() 184 s.inputPipe = inW 185 m.session[ref] = s 186 m.sessionMu.Unlock() 187 defer func() { 188 close(statusChan) 189 m.sessionMu.Lock() 190 s, ok := m.session[ref] 191 if ok { 192 s.statusChan = nil 193 s.buildOnGoing.Store(false) 194 } 195 m.sessionMu.Unlock() 196 }() 197 198 pw := pb.NewProgressWriter(statusChan) 199 200 // Build the specified request 201 ctx, cancel := context.WithCancel(ctx) 202 defer cancel() 203 resp, res, buildErr := m.buildFunc(ctx, req.Options, inR, pw) 204 m.sessionMu.Lock() 205 if s, ok := m.session[ref]; ok { 206 // NOTE: buildFunc can return *build.ResultHandle even on error (e.g. when it's implemented using (github.com/docker/buildx/controller/build).RunBuild). 207 if res != nil { 208 s.result = res 209 s.cancelBuild = cancel 210 s.buildOptions = req.Options 211 m.session[ref] = s 212 if buildErr != nil { 213 buildErr = controllererrors.WrapBuild(buildErr, ref) 214 } 215 } 216 } else { 217 m.sessionMu.Unlock() 218 return nil, errors.Errorf("build: unknown key %v", ref) 219 } 220 m.sessionMu.Unlock() 221 222 if buildErr != nil { 223 return nil, buildErr 224 } 225 226 if resp == nil { 227 resp = &client.SolveResponse{} 228 } 229 return &pb.BuildResponse{ 230 ExporterResponse: resp.ExporterResponse, 231 }, nil 232 } 233 234 func (m *Server) Status(req *pb.StatusRequest, stream pb.Controller_StatusServer) error { 235 ref := req.Ref 236 if ref == "" { 237 return errors.New("status: empty key") 238 } 239 240 // Wait and get status channel prepared by Build() 241 var statusChan <-chan *pb.StatusResponse 242 for { 243 // TODO: timeout? 244 m.sessionMu.Lock() 245 if _, ok := m.session[ref]; !ok || m.session[ref].statusChan == nil { 246 m.sessionMu.Unlock() 247 time.Sleep(time.Millisecond) // TODO: wait Build without busy loop and make it cancellable 248 continue 249 } 250 statusChan = m.session[ref].statusChan 251 m.sessionMu.Unlock() 252 break 253 } 254 255 // forward status 256 for ss := range statusChan { 257 if ss == nil { 258 break 259 } 260 if err := stream.Send(ss); err != nil { 261 return err 262 } 263 } 264 265 return nil 266 } 267 268 func (m *Server) Input(stream pb.Controller_InputServer) (err error) { 269 // Get the target ref from init message 270 msg, err := stream.Recv() 271 if err != nil { 272 if !errors.Is(err, io.EOF) { 273 return err 274 } 275 return nil 276 } 277 init := msg.GetInit() 278 if init == nil { 279 return errors.Errorf("unexpected message: %T; wanted init", msg.GetInit()) 280 } 281 ref := init.Ref 282 if ref == "" { 283 return errors.New("input: no ref is provided") 284 } 285 286 // Wait and get input stream pipe prepared by Build() 287 var inputPipeW *io.PipeWriter 288 for { 289 // TODO: timeout? 290 m.sessionMu.Lock() 291 if _, ok := m.session[ref]; !ok || m.session[ref].inputPipe == nil { 292 m.sessionMu.Unlock() 293 time.Sleep(time.Millisecond) // TODO: wait Build without busy loop and make it cancellable 294 continue 295 } 296 inputPipeW = m.session[ref].inputPipe 297 m.sessionMu.Unlock() 298 break 299 } 300 301 // Forward input stream 302 eg, ctx := errgroup.WithContext(context.TODO()) 303 done := make(chan struct{}) 304 msgCh := make(chan *pb.InputMessage) 305 eg.Go(func() error { 306 defer close(msgCh) 307 for { 308 msg, err := stream.Recv() 309 if err != nil { 310 if !errors.Is(err, io.EOF) { 311 return err 312 } 313 return nil 314 } 315 select { 316 case msgCh <- msg: 317 case <-done: 318 return nil 319 case <-ctx.Done(): 320 return nil 321 } 322 } 323 }) 324 eg.Go(func() (retErr error) { 325 defer close(done) 326 defer func() { 327 if retErr != nil { 328 inputPipeW.CloseWithError(retErr) 329 return 330 } 331 inputPipeW.Close() 332 }() 333 for { 334 var msg *pb.InputMessage 335 select { 336 case msg = <-msgCh: 337 case <-ctx.Done(): 338 return errors.Wrap(ctx.Err(), "canceled") 339 } 340 if msg == nil { 341 return nil 342 } 343 if data := msg.GetData(); data != nil { 344 if len(data.Data) > 0 { 345 _, err := inputPipeW.Write(data.Data) 346 if err != nil { 347 return err 348 } 349 } 350 if data.EOF { 351 return nil 352 } 353 } 354 } 355 }) 356 357 return eg.Wait() 358 } 359 360 func (m *Server) Invoke(srv pb.Controller_InvokeServer) error { 361 containerIn, containerOut := ioset.Pipe() 362 defer func() { containerOut.Close(); containerIn.Close() }() 363 364 initDoneCh := make(chan *processes.Process) 365 initErrCh := make(chan error) 366 eg, egCtx := errgroup.WithContext(context.TODO()) 367 srvIOCtx, srvIOCancel := context.WithCancel(egCtx) 368 eg.Go(func() error { 369 defer srvIOCancel() 370 return serveIO(srvIOCtx, srv, func(initMessage *pb.InitMessage) (retErr error) { 371 defer func() { 372 if retErr != nil { 373 initErrCh <- retErr 374 } 375 }() 376 ref := initMessage.Ref 377 cfg := initMessage.InvokeConfig 378 379 m.sessionMu.Lock() 380 s, ok := m.session[ref] 381 if !ok { 382 m.sessionMu.Unlock() 383 return errors.Errorf("invoke: unknown key %v", ref) 384 } 385 m.sessionMu.Unlock() 386 387 pid := initMessage.ProcessID 388 if pid == "" { 389 return errors.Errorf("invoke: specify process ID") 390 } 391 proc, ok := s.processes.Get(pid) 392 if !ok { 393 // Start a new process. 394 if cfg == nil { 395 return errors.New("no container config is provided") 396 } 397 var err error 398 proc, err = s.processes.StartProcess(pid, s.result, cfg) 399 if err != nil { 400 return err 401 } 402 } 403 // Attach containerIn to this process 404 proc.ForwardIO(&containerIn, srvIOCancel) 405 initDoneCh <- proc 406 return nil 407 }, &ioServerConfig{ 408 stdin: containerOut.Stdin, 409 stdout: containerOut.Stdout, 410 stderr: containerOut.Stderr, 411 // TODO: signal, resize 412 }) 413 }) 414 eg.Go(func() (rErr error) { 415 defer srvIOCancel() 416 // Wait for init done 417 var proc *processes.Process 418 select { 419 case p := <-initDoneCh: 420 proc = p 421 case err := <-initErrCh: 422 return err 423 case <-egCtx.Done(): 424 return egCtx.Err() 425 } 426 427 // Wait for IO done 428 select { 429 case <-srvIOCtx.Done(): 430 return srvIOCtx.Err() 431 case err := <-proc.Done(): 432 return err 433 case <-egCtx.Done(): 434 return egCtx.Err() 435 } 436 }) 437 438 return eg.Wait() 439 }