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  }