func main() {
///// ARGUMENT PARSING //////////////////////////////////////////////////////
flag.Parse()
if *version {
fmt.Printf("clique-agent v%s\n", clique.CliqueAgentVersion)
os.Exit(0)
}
level := logrus.InfoLevel
if *debug {
level = logrus.DebugLevel
}
if *configPath == "" {
fmt.Fprintf(os.Stderr, "`-config` option is required\n")
os.Exit(1)
}
///// LOGGING ///////////////////////////////////////////////////////////////
logger := &logrus.Logger{
Out: os.Stdout,
Level: level,
Formatter: new(logrus.TextFormatter),
}
logger.Debug("Initializing internals...")
///// CONFIGURATION /////////////////////////////////////////////////////////
cfg, err := config.NewConfig(*configPath)
if err != nil {
logger.Fatal(err.Error())
}
///// TRANSFER //////////////////////////////////////////////////////////////
// Protocol
t, err := setupTransferrer(logger, cfg)
if err != nil {
logger.Fatal(err.Error())
}
// Server
transferListener, err := net.Listen(
"tcp", fmt.Sprintf("0.0.0.0:%d", cfg.TransferPort),
)
if err != nil {
logger.Fatalf("Setting up transfer server: %s", err.Error())
}
transferServer := transfer.NewServer(
logger, transferListener, t.transferReceiver,
)
// Client
transferConnector := transfer.NewConnector()
transferClient := transfer.NewClient(
logger, transferConnector, t.transferSender,
)
///// SCHEDULING ////////////////////////////////////////////////////////////
schedRandGen := scheduler.NewCryptoUIG()
schedTaskSelector := &scheduler.LotteryTaskSelector{
Rand: schedRandGen,
}
schedClock := clock.NewClock()
sched := scheduler.NewScheduler(
logger,
schedTaskSelector, // scheduling algorithm
time.Second, // sleep between tasks
schedClock,
)
///// TRANSFER REGISTRY /////////////////////////////////////////////////////
transferRegistry := registry.NewRegistry()
///// DISPATCHER ////////////////////////////////////////////////////////////
dsptchr := &dispatcher.Dispatcher{
Scheduler: sched,
TransferInterruptible: t.interruptible,
TransferClient: transferClient,
ApiRegistry: transferRegistry,
Logger: logger,
}
///// API ///////////////////////////////////////////////////////////////////
var apiServer *api.Server
if cfg.APIPort != 0 {
apiServer = api.NewServer(
cfg.APIPort,
transferRegistry,
dsptchr,
)
}
///// SIGNAL HANDLER ////////////////////////////////////////////////////////
sigTermCh := make(chan os.Signal)
signal.Notify(sigTermCh, os.Interrupt)
signal.Notify(sigTermCh, syscall.SIGTERM)
go func() {
<-sigTermCh
logger.Debug("Closing transfer server...")
transferListener.Close()
logger.Debug("Closing scheduler...")
sched.Stop()
if apiServer != nil {
logger.Debug("Closing API server...")
apiServer.Close()
}
logger.Info("Exitting clique-agent...")
}()
logger.Debug("Initialization is complete!")
///// START /////////////////////////////////////////////////////////////////
logger.Info("Clique Agent")
// Populate the dispatcher with tasks
if len(cfg.RemoteHosts) != 0 {
createTransferTasks(logger, cfg, dsptchr)
}
wg := new(sync.WaitGroup)
// Start the transfer server
wg.Add(1)
go func() {
transferServer.Serve()
logger.Debug("Transfer server is done.")
wg.Done()
}()
// Start the scheduler
wg.Add(1)
go func() {
sched.Run()
logger.Debug("Scheduler is done.")
wg.Done()
}()
// Start the API server
if apiServer != nil {
wg.Add(1)
go func() {
apiServer.Serve()
logger.Debug("API server is done.")
wg.Done()
}()
}
// Wait until everything is done!
wg.Wait()
logger.Debug("Clique agent is done.")
}
server *transfer.Server
listenerConnChan chan net.Conn
listenerErrChan chan error
serverClosed chan struct{}
)
BeforeEach(func() {
logger = &logrus.Logger{
Out: GinkgoWriter,
Level: logrus.DebugLevel,
Formatter: new(logrus.TextFormatter),
}
fakeListener = new(fakes.FakeListener)
fakeTransferReceiver = new(fakes.FakeTransferReceiver)
server = transfer.NewServer(logger, fakeListener, fakeTransferReceiver)
listenerConnChan = make(chan net.Conn, 100)
listenerErrChan = make(chan error)
fakeListener.AcceptStub = func() (net.Conn, error) {
select {
case conn := <-listenerConnChan:
return conn, nil
case err := <-listenerErrChan:
return nil, err
}
}
})
JustBeforeEach(func() {
serverClosed = make(chan struct{})