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.")
}
)
BeforeEach(func() {
logger = &logrus.Logger{
Out: GinkgoWriter,
Level: logrus.DebugLevel,
Formatter: new(logrus.TextFormatter),
}
fakeConnector = new(fakes.FakeConnector)
conn, _ = net.Pipe()
fakeConnector.ConnectReturns(conn, nil)
fakeTransferSender = new(fakes.FakeTransferSender)
client = transfer.NewClient(logger, fakeConnector, fakeTransferSender)
})
It("should create a connection to the server", func() {
spec := transfer.TransferSpec{
IP: net.ParseIP("127.0.0.1"),
Port: 1200,
}
_, err := client.Transfer(spec)
Expect(err).NotTo(HaveOccurred())
Expect(fakeConnector.ConnectCallCount()).To(Equal(1))
receivedIP, receivedPort := fakeConnector.ConnectArgsForCall(0)
Expect(receivedIP).To(Equal(spec.IP))
Expect(receivedPort).To(Equal(spec.Port))
})