func (s *WardenServer) streamProcessToConnection(processID uint32, stream <-chan warden.ProcessStream, conn net.Conn) proto.Message {
for payload := range stream {
if payload.ExitStatus != nil {
return &protocol.ProcessPayload{
ProcessId: proto.Uint32(processID),
ExitStatus: proto.Uint32(*payload.ExitStatus),
}
}
var payloadSource protocol.ProcessPayload_Source
switch payload.Source {
case warden.ProcessStreamSourceStdout:
payloadSource = protocol.ProcessPayload_stdout
case warden.ProcessStreamSourceStderr:
payloadSource = protocol.ProcessPayload_stderr
case warden.ProcessStreamSourceStdin:
payloadSource = protocol.ProcessPayload_stdin
}
protocol.Messages(&protocol.ProcessPayload{
ProcessId: proto.Uint32(processID),
Source: &payloadSource,
Data: proto.String(string(payload.Data)),
}).WriteTo(conn)
}
return nil
}
func (s *WardenServer) handleRun(conn net.Conn, request *protocol.RunRequest) (proto.Message, error) {
handle := request.GetHandle()
script := request.GetScript()
privileged := request.GetPrivileged()
env := request.GetEnv()
container, err := s.backend.Lookup(handle)
if err != nil {
return nil, err
}
s.bomberman.Pause(container.Handle())
defer s.bomberman.Unpause(container.Handle())
ProcessSpec := warden.ProcessSpec{
Script: script,
Privileged: privileged,
EnvironmentVariables: convertEnvironmentVariables(env),
}
if request.Rlimits != nil {
ProcessSpec.Limits = resourceLimits(request.Rlimits)
}
processID, stream, err := container.Run(ProcessSpec)
if err != nil {
return nil, err
}
protocol.Messages(&protocol.ProcessPayload{
ProcessId: proto.Uint32(processID),
}).WriteTo(conn)
return s.streamProcessToConnection(processID, stream, conn), nil
}
func (s *WardenServer) handleSpawn(request *protocol.SpawnRequest) (proto.Message, error) {
handle := request.GetHandle()
script := request.GetScript()
privileged := request.GetPrivileged()
discardOutput := request.GetDiscardOutput()
container, err := s.backend.Lookup(handle)
if err != nil {
return nil, err
}
s.bomberman.Pause(container.Handle())
defer s.bomberman.Unpause(container.Handle())
jobSpec := backend.JobSpec{
Script: script,
Privileged: privileged,
DiscardOutput: discardOutput,
AutoLink: true,
}
if request.Rlimits != nil {
jobSpec.Limits = resourceLimits(request.Rlimits)
}
jobID, err := container.Spawn(jobSpec)
if err != nil {
return nil, err
}
return &protocol.SpawnResponse{JobId: proto.Uint32(jobID)}, nil
}
func (s *WardenServer) handleNetIn(request *protocol.NetInRequest) (proto.Message, error) {
handle := request.GetHandle()
hostPort := request.GetHostPort()
containerPort := request.GetContainerPort()
container, err := s.backend.Lookup(handle)
if err != nil {
return nil, err
}
s.bomberman.Pause(container.Handle())
defer s.bomberman.Unpause(container.Handle())
hostPort, containerPort, err = container.NetIn(hostPort, containerPort)
if err != nil {
return nil, err
}
return &protocol.NetInResponse{
HostPort: proto.Uint32(hostPort),
ContainerPort: proto.Uint32(containerPort),
}, nil
}
func (c *Connection) Link(handle string, jobID uint32) (*warden.LinkResponse, error) {
res, err := c.RoundTrip(
&warden.LinkRequest{
Handle: proto.String(handle),
JobId: proto.Uint32(jobID),
},
&warden.LinkResponse{},
)
if err != nil {
return nil, err
}
return res.(*warden.LinkResponse), nil
}
func (f *FakeGordon) Run(handle string, script string) (*warden.RunResponse, error) {
f.lock.Lock()
defer f.lock.Unlock()
f.scriptsThatRan = append(f.scriptsThatRan, &RunningScript{
Handle: handle,
Script: script,
})
for ro, cb := range f.runCallbacks {
if ro.Handle == handle && ro.Script == script {
return cb()
}
}
return &warden.RunResponse{ExitStatus: proto.Uint32(f.runReturnStatusCode)}, f.runReturnError
}
func (s *WardenServer) handleRun(request *protocol.RunRequest) (proto.Message, error) {
handle := request.GetHandle()
script := request.GetScript()
privileged := request.GetPrivileged()
discardOutput := request.GetDiscardOutput()
container, err := s.backend.Lookup(handle)
if err != nil {
return nil, err
}
s.bomberman.Pause(container.Handle())
defer s.bomberman.Unpause(container.Handle())
jobSpec := backend.JobSpec{
Script: script,
Privileged: privileged,
DiscardOutput: discardOutput,
AutoLink: false,
}
if request.Rlimits != nil {
jobSpec.Limits = resourceLimits(request.Rlimits)
}
jobID, err := container.Spawn(jobSpec)
if err != nil {
return nil, err
}
jobResult, err := container.Link(jobID)
if err != nil {
return nil, err
}
return &protocol.RunResponse{
ExitStatus: proto.Uint32(jobResult.ExitStatus),
Stdout: proto.String(string(jobResult.Stdout)),
Stderr: proto.String(string(jobResult.Stderr)),
}, nil
}
func (c *Connection) Stream(handle string, jobId uint32) (chan *warden.StreamResponse, chan bool, error) {
err := c.sendMessage(
&warden.StreamRequest{
Handle: proto.String(handle),
JobId: proto.Uint32(jobId),
},
)
if err != nil {
return nil, nil, err
}
responses := make(chan *warden.StreamResponse)
streamDone := make(chan bool)
go func() {
for {
resMsg, err := c.readResponse(&warden.StreamResponse{})
if err != nil {
close(responses)
close(streamDone)
break
}
response := resMsg.(*warden.StreamResponse)
responses <- response
if response.ExitStatus != nil {
close(responses)
close(streamDone)
break
}
}
}()
return responses, streamDone, nil
}
func (s *WardenServer) handleLink(link *protocol.LinkRequest) (proto.Message, error) {
handle := link.GetHandle()
jobID := link.GetJobId()
container, err := s.backend.Lookup(handle)
if err != nil {
return nil, err
}
s.bomberman.Pause(container.Handle())
defer s.bomberman.Unpause(container.Handle())
jobResult, err := container.Link(jobID)
if err != nil {
return nil, err
}
return &protocol.LinkResponse{
ExitStatus: proto.Uint32(jobResult.ExitStatus),
Stdout: proto.String(string(jobResult.Stdout)),
Stderr: proto.String(string(jobResult.Stderr)),
}, nil
}
func (s *WardenServer) handleStream(conn net.Conn, request *protocol.StreamRequest) (proto.Message, error) {
handle := request.GetHandle()
jobID := request.GetJobId()
container, err := s.backend.Lookup(handle)
if err != nil {
return nil, err
}
s.bomberman.Pause(container.Handle())
defer s.bomberman.Unpause(container.Handle())
stream, err := container.Stream(jobID)
if err != nil {
return nil, err
}
var response proto.Message
for chunk := range stream {
if chunk.ExitStatus != nil {
response = &protocol.StreamResponse{
ExitStatus: proto.Uint32(*chunk.ExitStatus),
}
break
}
protocol.Messages(&protocol.StreamResponse{
Name: proto.String(chunk.Name),
Data: proto.String(string(chunk.Data)),
}).WriteTo(conn)
}
return response, nil
}
Properties: map[string]string{
"foo": "bar",
},
})
Ω(err).ShouldNot(HaveOccurred())
Ω(handle).Should(Equal("foohandle"))
ro := protocol.CreateRequest_BindMount_RO
rw := protocol.CreateRequest_BindMount_RW
hostOrigin := protocol.CreateRequest_BindMount_Host
containerOrigin := protocol.CreateRequest_BindMount_Container
assertWriteBufferContains(&protocol.CreateRequest{
Handle: proto.String("some-handle"),
GraceTime: proto.Uint32(10),
Rootfs: proto.String("some-rootfs-path"),
Network: proto.String("some-network"),
BindMounts: []*protocol.CreateRequest_BindMount{
{
SrcPath: proto.String("/src-a"),
DstPath: proto.String("/dst-a"),
Mode: &ro,
Origin: &hostOrigin,
},
{
SrcPath: proto.String("/src-b"),
DstPath: proto.String("/dst-b"),
Mode: &rw,
Origin: &containerOrigin,
},
// InitChecksum initializes a checksum based on the provided key and
// the contents of the value. If the value contains a byte slice, the
// checksum includes it directly; if the value contains an integer,
// the checksum includes the integer as 8 bytes in big-endian order.
func (v *Value) InitChecksum(key []byte) {
if v.Checksum == nil {
v.Checksum = gogoproto.Uint32(v.computeChecksum(key))
}
}
wardenServer.Stop()
Expect(time.Since(before)).To(BeNumerically(">", 50*time.Millisecond))
readResponse(&protocol.CreateResponse{})
_, err := protocol.Messages(&protocol.PingRequest{}).WriteTo(serverConnection)
Expect(err).To(HaveOccurred())
})
})
dontWaitRequests := []proto.Message{
&protocol.LinkRequest{
Handle: proto.String("some-handle"),
JobId: proto.Uint32(1),
},
&protocol.StreamRequest{
Handle: proto.String("some-handle"),
JobId: proto.Uint32(1),
},
&protocol.RunRequest{
Handle: proto.String("some-handle"),
Script: proto.String("some-script"),
},
}
for _, req := range dontWaitRequests {
request := req
Context(fmt.Sprintf("when a %T request is in-flight", request), func() {
readResponse(&protocol.CreateResponse{})
_, err := protocol.Messages(&protocol.PingRequest{}).WriteTo(serverConnection)
Expect(err).To(HaveOccurred())
})
})
dontWaitRequests := []proto.Message{
&protocol.RunRequest{
Handle: proto.String("some-handle"),
Script: proto.String("some-script"),
},
&protocol.AttachRequest{
Handle: proto.String("some-handle"),
ProcessId: proto.Uint32(1),
},
}
for _, req := range dontWaitRequests {
request := req
Context(fmt.Sprintf("when a %T request is in-flight", request), func() {
BeforeEach(func() {
serverBackend = fake_backend.NewSlow(100 * time.Millisecond)
container, err := serverBackend.Create(warden.ContainerSpec{Handle: "some-handle"})
Expect(err).ToNot(HaveOccurred())
exitStatus := uint32(42)
Background: proto.Bool(true),
Kill: proto.Bool(true),
},
&warden.DestroyRequest{Handle: proto.String("foo")},
).Bytes())
Ω(string(writeBuffer.Bytes())).Should(Equal(expectedWriteBufferContents))
})
})
Describe("Spawning and streaming", func() {
BeforeEach(func() {
provider = NewFakeConnectionProvider(
warden.Messages(
&warden.SpawnResponse{
JobId: proto.Uint32(42),
},
&warden.StreamResponse{
Name: proto.String("stdout"),
Data: proto.String("some data for stdout"),
},
),
writeBuffer,
)
client = NewClient(provider)
err := client.Connect()
Ω(err).ShouldNot(HaveOccurred())
})
It("should spawn and stream succesfully", func(done Done) {
It("should return the correct memory format", func() {
diskLimit, err := connection.GetDiskLimit("foo")
Ω(err).ShouldNot(HaveOccurred())
Ω(diskLimit).Should(BeNumerically("==", 40))
assertWriteBufferContains(&warden.LimitDiskRequest{
Handle: proto.String("foo"),
})
})
})
})
Describe("Spawning", func() {
BeforeEach(func() {
wardenMessages = append(wardenMessages,
&warden.SpawnResponse{JobId: proto.Uint32(42)},
&warden.SpawnResponse{JobId: proto.Uint32(43)},
)
})
It("should be able to spawn multiple jobs sequentially", func() {
resp, err := connection.Spawn("foo-handle", "echo hi", true)
Ω(err).ShouldNot(HaveOccurred())
Ω(resp.GetJobId()).Should(BeNumerically("==", 42))
assertWriteBufferContains(&warden.SpawnRequest{
Handle: proto.String("foo-handle"),
Script: proto.String("echo hi"),
DiscardOutput: proto.Bool(true),
})