From 8bfdce5ec1053edcf9f60086a12800c6fc532f86 Mon Sep 17 00:00:00 2001
From: nbari <nbari@tequila.io>
Date: Wed, 8 Nov 2017 14:10:19 +0100
Subject: [PATCH] use sync.Map

---
 .travis.yml  |  8 +++++--
 scheduler.go | 62 +++++++++++++++++++++-------------------------------
 2 files changed, 31 insertions(+), 39 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 9dc6a46..b063e1f 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -1,8 +1,12 @@
 language: go
 
+os:
+  - linux
+  - osx
+
 go:
-  - 1.7
-  - tip
+  - 1.9.*
+  - master
 
 before_install:
   - go get github.com/axw/gocov/gocov
diff --git a/scheduler.go b/scheduler.go
index 636d854..443110d 100644
--- a/scheduler.go
+++ b/scheduler.go
@@ -7,80 +7,68 @@ import (
 	"time"
 )
 
-// Scheduler map of schedulers
-type Scheduler struct {
-	Schedulers map[string]scheduler
-	sync.Mutex
-}
-
+// scheduler private struct
 type scheduler struct {
 	t    <-chan time.Time
 	quit chan struct{}
 	f    func()
 }
 
+// Scheduler map of schedulers
+type Scheduler struct {
+	sync.Map
+}
+
 // New returns a new scheduler
 func New() *Scheduler {
-	return &Scheduler{
-		Schedulers: make(map[string]scheduler),
-	}
+	return &Scheduler{}
 }
 
 // AddScheduler calls a function every X seconds.
 func (s *Scheduler) AddScheduler(name string, interval int, f func()) {
-	s.Lock()
-	defer s.Unlock()
-
 	e := time.Duration(interval) * time.Second
 
-	scheduler := scheduler{
+	// create a new scheduler
+	task := scheduler{
 		t:    time.NewTicker(e).C,
 		quit: make(chan struct{}),
 		f:    f,
 	}
 
-	// stop scheduler if exist
-	if sk, ok := s.Schedulers[name]; ok {
-		close(sk.quit)
+	// stop scheduler if exist or add a new one
+	if sk, ok := s.LoadOrStore(name, task); ok {
+		close(sk.(scheduler).quit)
 	}
 
-	// add service
-	s.Schedulers[name] = scheduler
-
-	go func() {
+	// Create the scheduler in a goroutine running forever until it quits
+	go func(s scheduler) {
 		for {
 			select {
-			case <-scheduler.t:
-				scheduler.f()
-			case <-scheduler.quit:
+			case <-s.t:
+				s.f()
+			case <-s.quit:
 				return
 			}
 		}
-	}()
+	}(task)
 }
 
 // Stop ends a specified scheduler.
 func (s *Scheduler) Stop(name string) error {
-	s.Lock()
-	defer s.Unlock()
-
-	sk, ok := s.Schedulers[name]
-
+	sk, ok := s.Load(name)
 	if !ok {
 		return fmt.Errorf("Scheduler: %s, does not exist.", name)
 	}
-
-	close(sk.quit)
+	close(sk.(scheduler).quit)
 	return nil
 }
 
 // StopAll ends all schedulers.
 func (s *Scheduler) StopAll() {
-	s.Lock()
-	defer s.Unlock()
-
-	for k, v := range s.Schedulers {
-		close(v.quit)
-		log.Printf("Stoping: %s", k)
+	close := func(key, value interface{}) bool {
+		close(value.(scheduler).quit)
+		log.Printf("Stoping: %s", key)
+		return true
 	}
+	s.Range(close)
 }