func (r *AgentServer) init() (err error) {
tlsConfig := r.Option.CertFiles.MakeTLSConfig()
if tlsConfig == nil {
r.listener, err = net.Listen("tcp", *r.Option.Host+":"+strconv.Itoa(*r.Option.Port))
} else {
r.listener, err = tls.Listen("tcp", *r.Option.Host+":"+strconv.Itoa(*r.Option.Port), tlsConfig)
}
util.SetupHttpClient(tlsConfig)
if err != nil {
log.Fatal(err)
}
fmt.Println("AgentServer starts on", *r.Option.Host+":"+strconv.Itoa(*r.Option.Port))
if *r.Option.CleanRestart {
if fileInfos, err := ioutil.ReadDir(r.storageBackend.dir); err == nil {
suffix := fmt.Sprintf("-%d.dat", *r.Option.Port)
for _, fi := range fileInfos {
name := fi.Name()
if !fi.IsDir() && strings.HasSuffix(name, suffix) {
// println("removing old dat file:", name)
os.Remove(filepath.Join(r.storageBackend.dir, name))
}
}
}
}
return
}
func RunMaster(tlsConfig *tls.Config, listenOn string) {
tl := &TeamMaster{}
tl.channels = &NamedChannelMap{name2Chans: make(map[string][]*ChannelInformation)}
tl.MasterResource = NewMasterResource()
masterMux := http.NewServeMux()
masterMux.HandleFunc("/", tl.statusHandler)
masterMux.HandleFunc("/agent/assign", tl.requestAgentHandler)
masterMux.HandleFunc("/agent/update", tl.updateAgentHandler)
masterMux.HandleFunc("/agent/", tl.listAgentsHandler)
masterMux.HandleFunc("/channel/", tl.handleChannel)
var listener net.Listener
var err error
if tlsConfig == nil {
listener, err = net.Listen("tcp", listenOn)
} else {
listener, err = tls.Listen("tcp", listenOn, tlsConfig)
}
if err != nil {
log.Fatalf("Volume server fail to serve public: %v", err)
}
util.SetupHttpClient(tlsConfig)
if e := http.Serve(listener, masterMux); e != nil {
log.Fatalf("Volume server fail to serve public: %v", e)
}
}
// if this should not run, return false
func (tr *TaskRunner) Run(fc *flow.FlowContext) {
if fc.Id != tr.option.ContextId {
return
}
fc.ChannelBufferSize = tr.option.ChannelBufferSize
// using driverOption, not so clean way
tr.option.TaskTlsConfig = driverOption.CertFiles.MakeTLSConfig()
util.SetupHttpClient(tr.option.TaskTlsConfig)
_, taskGroups := plan.GroupTasks(fc)
tr.Tasks = taskGroups[tr.option.TaskGroupId].Tasks
tr.FlowContext = fc
tr.executorStatus.StartTime = time.Now()
go tr.reportLocalExecutorStatus()
// println("taskGroup", tr.Tasks[0].Name(), "starts")
// 4. setup task input and output channels
var wg sync.WaitGroup
tr.connectInputsAndOutputs(&wg)
// 6. starts to run the task locally
for _, task := range tr.Tasks {
// println("run task", task.Name())
wg.Add(1)
go func(task *flow.Task) {
defer wg.Done()
task.RunTask()
}(task)
}
// 7. need to close connected output channels
wg.Wait()
// println("taskGroup", tr.Tasks[0].Name(), "finishes", tr.option.RequestId)
tr.executorStatus.StopTime = time.Now()
tr.reportLocalExecutorStatusOnce()
}
func TestNormalHttpCopy(t *testing.T) {
util.SetupHttpClient(nil)
files := []string{
os.Args[0], "http_server.go", "../scheduler/scheduler.go", "fetch_url.go",
}
rsyncServer, err := NewRsyncServer(files[0], files[1:])
if err != nil {
t.Fatalf("Failed to start local server: %v", err)
}
rsyncServer.StartRsyncServer(nil, ":0")
err = FetchFilesTo(fmt.Sprintf("localhost:%d", rsyncServer.Port), "/tmp")
if err != nil {
fmt.Printf("pausing localhost:%d\n", rsyncServer.Port)
time.Sleep(time.Minute)
t.Fatalf("Failed to download file: %v", err)
}
t.Logf("file downloading works ok.")
}
// driver runs on local, controlling all tasks
func (fcd *FlowContextDriver) Run(fc *flow.FlowContext) {
fcd.checkParameters()
// task fusion to minimize disk IO
fcd.stepGroups, fcd.taskGroups = plan.GroupTasks(fc)
// plot the execution graph
if fcd.Option.PlotOutput {
plan.PlotGraph(fcd.taskGroups, fc)
return
}
tlsConfig := fcd.Option.CertFiles.MakeTLSConfig()
util.SetupHttpClient(tlsConfig)
// start server to serve files to agents to run exectuors
rsyncServer, err := rsync.NewRsyncServer(os.Args[0], fcd.Option.RelatedFileNames())
if err != nil {
log.Fatalf("Failed to start local server: %v", err)
}
rsyncServer.StartRsyncServer(tlsConfig, fcd.Option.Host+":"+strconv.Itoa(fcd.Option.Port))
driverHost := fcd.Option.Host
// create thes cheduler
sched := scheduler.NewScheduler(
fcd.Option.Leader,
&scheduler.SchedulerOption{
DataCenter: fcd.Option.DataCenter,
Rack: fcd.Option.Rack,
TaskMemoryMB: fcd.Option.TaskMemoryMB,
DriverHost: driverHost,
DriverPort: rsyncServer.Port,
Module: fcd.Option.Module,
ExecutableFile: os.Args[0],
ExecutableFileHash: rsyncServer.ExecutableFileHash(),
TlsConfig: tlsConfig,
},
)
// best effort to clean data on agent disk
// this may need more improvements
defer fcd.Cleanup(sched, fc)
go sched.EventLoop()
flow.OnInterrupt(func() {
fcd.OnInterrupt(fc, sched)
}, func() {
fcd.OnExit(fc, sched)
})
// schedule to run the steps
var wg sync.WaitGroup
for _, taskGroup := range fcd.taskGroups {
wg.Add(1)
sched.EventChan <- scheduler.SubmitTaskGroup{
FlowContext: fc,
TaskGroup: taskGroup,
Bid: fcd.Option.FlowBid / float64(len(fcd.taskGroups)),
WaitGroup: &wg,
}
}
go sched.Market.FetcherLoop()
wg.Wait()
fcd.CloseOutputChannels(fc)
if fcd.Option.ShowFlowStats {
fcd.ShowFlowStatus(fc, sched)
}
}
func main() {
master.Flag("cert.file", "A PEM eoncoded certificate file").Default("").StringVar(&masterCerts.CertFile)
master.Flag("key.file", "A PEM encoded private key file.").Default("").StringVar(&masterCerts.KeyFile)
master.Flag("ca.file", "A PEM eoncoded CA's certificate file.").Default("").StringVar(&masterCerts.CaFile)
agent.Flag("cert.file", "A PEM eoncoded certificate file").Default("").StringVar(&agentOption.CertFiles.CertFile)
agent.Flag("key.file", "A PEM encoded private key file.").Default("").StringVar(&agentOption.CertFiles.KeyFile)
agent.Flag("ca.file", "A PEM eoncoded CA's certificate file.").Default("").StringVar(&agentOption.CertFiles.CaFile)
sender.Flag("cert.file", "A PEM eoncoded certificate file").Default("").StringVar(&senderCerts.CertFile)
sender.Flag("key.file", "A PEM encoded private key file.").Default("").StringVar(&senderCerts.KeyFile)
sender.Flag("ca.file", "A PEM eoncoded CA's certificate file.").Default("").StringVar(&senderCerts.CaFile)
receiver.Flag("cert.file", "A PEM eoncoded certificate file").Default("").StringVar(&receiverCerts.CertFile)
receiver.Flag("key.file", "A PEM encoded private key file.").Default("").StringVar(&receiverCerts.KeyFile)
receiver.Flag("ca.file", "A PEM eoncoded CA's certificate file.").Default("").StringVar(&receiverCerts.CaFile)
switch kingpin.MustParse(app.Parse(os.Args[1:])) {
case master.FullCommand():
println("listening on", *masterAddress)
m.RunMaster(masterCerts.MakeTLSConfig(), *masterAddress)
case sender.FullCommand():
tlsConfig := senderCerts.MakeTLSConfig()
util.SetupHttpClient(tlsConfig)
var wg sync.WaitGroup
sendChan, err := s.NewDirectSendChannel(tlsConfig, *sendToChanName, *senderAgentAddress, &wg)
if err != nil {
panic(err)
}
file := os.Stdin
if *sendFile != "" {
file, err = os.Open(*sendFile)
if err != nil {
log.Fatal(err)
}
defer file.Close()
}
counter := 0
scanner := bufio.NewScanner(file)
for scanner.Scan() {
sendChan <- scanner.Bytes()
counter++
}
if err := scanner.Err(); err != nil {
log.Fatal(err)
}
close(sendChan)
wg.Wait()
case receiver.FullCommand():
tlsConfig := receiverCerts.MakeTLSConfig()
util.SetupHttpClient(tlsConfig)
target := r.FindTarget(*receiveFromChanName, *receiverMaster)
rc := r.NewReceiveChannel(tlsConfig, *receiveFromChanName, 0)
recvChan, err := rc.GetDirectChannel(target, 128)
if err != nil {
panic(err)
}
for m := range recvChan {
println(string(m))
}
case agent.FullCommand():
if agentOption.CertFiles.IsEnabled() {
if *agentOption.Host == "" {
log.Fatalf("Usage Note: --host option is needed in 2-way SSL mode and must match CN in the certificate.")
}
}
agentServer := a.NewAgentServer(agentOption)
agentServer.Run()
}
}