implement the ChannelSource

Author: lichuqiang

pkg/controller/manager.go

@@ -146,6 +146,9 @@ func (cm *controllerManager) injectInto(i interface{}) error {
 	if _, err := inject.CacheInto(cm.cache, i); err != nil {
 		return err
 	}
+	if _, err := inject.DoStop(cm.stop, i); err != nil {
+		return err
+	}
 	return nil
 }
 

pkg/controller/source/example_test.go

@@ -41,7 +41,7 @@ func ExampleChannelSource() {
 	events := make(chan event.GenericEvent)
 
 	ctrl.Watch(
-		source.ChannelSource(events),
+		&source.ChannelSource{Source: events},
 		&eventhandler.EnqueueHandler{},
 	)
 }

pkg/controller/source/source.go

@@ -18,6 +18,7 @@ package source
 
 import (
 	"fmt"
+	"sync"
 
 	"github.com/kubernetes-sigs/controller-runtime/pkg/controller/event"
 	"github.com/kubernetes-sigs/controller-runtime/pkg/controller/eventhandler"
@@ -30,6 +31,11 @@ import (
 	"github.com/kubernetes-sigs/controller-runtime/pkg/controller/predicate"
 )
 
+const (
+	// defaultBufferSize is the default number of event notifications that can be buffered.
+	defaultBufferSize = 1024
+)
+
 // Source is a source of events (eh.g. Create, Update, Delete operations on Kubernetes Objects, Webhook callbacks, etc)
 // which should be processed by event.EventHandlers to enqueue reconcile.Requests.
 //
@@ -45,21 +51,127 @@ type Source interface {
 	Start(eventhandler.EventHandler, workqueue.RateLimitingInterface, ...predicate.Predicate) error
 }
 
+var _ Source = &ChannelSource{}
+
 // ChannelSource is used to provide a source of events originating outside the cluster
 // (eh.g. GitHub Webhook callback).  ChannelSource requires the user to wire the external
 // source (eh.g. http handler) to write GenericEvents to the underlying channel.
-type ChannelSource chan event.GenericEvent
+type ChannelSource struct {
+	// once ensures the event distribution goroutine will be performed only once
+	once sync.Once
+
+	// Source is the source channel to fetch GenericEvents
+	Source <-chan event.GenericEvent
+
+	// stop is to end ongoing goroutine, and close the channels
+	stop <-chan struct{}
+
+	// dest is the destination channels of the added event handlers
+	dest []chan event.GenericEvent
+
+	// DestBufferSize is the specified buffer size of dest channels
+	DestBufferSize int
 
-var _ Source = ChannelSource(make(chan event.GenericEvent))
+	// destLock is to ensure the destination channels are safely added/removed
+	destLock sync.Mutex
+}
+
+var _ inject.Stop = &ChannelSource{}
+
+// InjectStop is internal should be called only by the Controller.
+// It is used to inject the stop channel initialized by the ControllerManager.
+func (cs *ChannelSource) InjectStop(stop <-chan struct{}) error {
+	if cs.stop == nil {
+		cs.stop = stop
+	}
+
+	return nil
+}
 
 // Start implements Source and should only be called by the Controller.
-func (ks ChannelSource) Start(
+func (cs *ChannelSource) Start(
 	handler eventhandler.EventHandler,
 	queue workqueue.RateLimitingInterface,
 	prct ...predicate.Predicate) error {
+	// Source should have been specified by the user.
+	if cs.Source == nil {
+		return fmt.Errorf("must specify ChannelSource.Source")
+	}
+
+	// stop should have been injected before Start was called
+	if cs.stop == nil {
+		return fmt.Errorf("must call InjectStop on ChannelSource before calling Start")
+	}
+
+	// use default value if DestBufferSize not specified
+	if cs.DestBufferSize == 0 {
+		cs.DestBufferSize = defaultBufferSize
+	}
+
+	cs.once.Do(func() {
+		// Distribute GenericEvents to all EventHandler / Queue pairs Watching this source
+		go func() {
+			for {
+				select {
+				case <-cs.stop:
+					// Close destination channels
+					cs.doStop()
+					return
+				case evt := <-cs.Source:
+					cs.distribute(evt)
+				}
+			}
+		}()
+	})
+
+	dst := make(chan event.GenericEvent, cs.DestBufferSize)
+	go func() {
+		for evt := range dst {
+			shouldHandle := true
+			for _, p := range prct {
+				if !p.Generic(evt) {
+					shouldHandle = false
+					break
+				}
+			}
+
+			if shouldHandle {
+				handler.Generic(queue, evt)
+			}
+		}
+	}()
+
+	cs.destLock.Lock()
+	defer cs.destLock.Unlock()
+
+	cs.dest = append(cs.dest, dst)
+
 	return nil
 }
 
+func (cs *ChannelSource) doStop() {
+	cs.destLock.Lock()
+	defer cs.destLock.Unlock()
+
+	for _, dst := range cs.dest {
+		close(dst)
+	}
+}
+
+func (cs *ChannelSource) distribute(evt event.GenericEvent) {
+	cs.destLock.Lock()
+	defer cs.destLock.Unlock()
+
+	for _, dst := range cs.dest {
+		// We cannot make it under goroutine here, or we'll meet the
+		// race condition of writing message to closed channels.
+		// To avoid blocking, the dest channels are expected to be of
+		// proper buffer size. If we still see it blocked, then
+		// the controller is thought to be in an abnormal state.
+		dst <- evt
+	}
+}
+
 // KindSource is used to provide a source of events originating inside the cluster from Watches (eh.g. Pod Create)
 type KindSource struct {
 	// Type is the type of object to watch.  e.g. &v1.Pod{}

pkg/controller/source/source_test.go

@@ -18,6 +18,7 @@ package source_test
 
 import (
 	"fmt"
+	"time"
 
 	"github.com/kubernetes-sigs/controller-runtime/pkg/cache/informertest"
 	"github.com/kubernetes-sigs/controller-runtime/pkg/controller/event"
@@ -26,10 +27,11 @@ import (
 	"github.com/kubernetes-sigs/controller-runtime/pkg/runtime/inject"
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
-	"k8s.io/client-go/util/workqueue"
 
 	"github.com/kubernetes-sigs/controller-runtime/pkg/controller/predicate"
 	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/client-go/util/workqueue"
 )
 
 var _ = Describe("Source", func() {
@@ -251,11 +253,194 @@ var _ = Describe("Source", func() {
 	})
 
 	Describe("ChannelSource", func() {
-		It("TODO(community): implement this", func(done Done) {
-			instance := source.ChannelSource(make(chan event.GenericEvent))
-			instance.Start(nil, nil)
+		var stop chan struct{}
+		var ch chan event.GenericEvent
 
-			close(done)
+		BeforeEach(func() {
+			stop = make(chan struct{})
+			ch = make(chan event.GenericEvent)
+		})
+
+		AfterEach(func() {
+			close(stop)
+			close(ch)
+		})
+
+		Context("for a source", func() {
+			It("should provide a GenericEvent", func(done Done) {
+				ch := make(chan event.GenericEvent)
+				c := make(chan struct{})
+				p := &corev1.Pod{
+					ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "bar"},
+				}
+				evt := event.GenericEvent{
+					Object: p,
+					Meta:   p,
+				}
+
+				q := workqueue.NewNamedRateLimitingQueue(workqueue.DefaultControllerRateLimiter(), "test")
+				instance := &source.ChannelSource{Source: ch}
+				inject.DoStop(stop, instance)
+				err := instance.Start(eventhandler.Funcs{
+					CreateFunc: func(workqueue.RateLimitingInterface, event.CreateEvent) {
+						defer GinkgoRecover()
+						Fail("Unexpected CreateEvent")
+					},
+					UpdateFunc: func(workqueue.RateLimitingInterface, event.UpdateEvent) {
+						defer GinkgoRecover()
+						Fail("Unexpected UpdateEvent")
+					},
+					DeleteFunc: func(workqueue.RateLimitingInterface, event.DeleteEvent) {
+						defer GinkgoRecover()
+						Fail("Unexpected DeleteEvent")
+					},
+					GenericFunc: func(q2 workqueue.RateLimitingInterface, evt event.GenericEvent) {
+						defer GinkgoRecover()
+						Expect(q2).To(Equal(q))
+						Expect(evt.Meta).To(Equal(p))
+						Expect(evt.Object).To(Equal(p))
+						close(c)
+					},
+				}, q)
+				Expect(err).NotTo(HaveOccurred())
+
+				ch <- evt
+				<-c
+				close(done)
+			})
+			It("should block if exceed buffer size", func(done Done) {
+				ch := make(chan event.GenericEvent)
+				evt := event.GenericEvent{}
+				interval := 5 * time.Second
+
+				q := workqueue.NewNamedRateLimitingQueue(workqueue.DefaultControllerRateLimiter(), "test")
+				// Add a handler to get distribution blocked
+				instance := &source.ChannelSource{Source: ch}
+				instance.DestBufferSize = 1
+				inject.DoStop(stop, instance)
+				err := instance.Start(eventhandler.Funcs{
+					CreateFunc: func(workqueue.RateLimitingInterface, event.CreateEvent) {
+						defer GinkgoRecover()
+						Fail("Unexpected CreateEvent")
+					},
+					UpdateFunc: func(workqueue.RateLimitingInterface, event.UpdateEvent) {
+						defer GinkgoRecover()
+						Fail("Unexpected UpdateEvent")
+					},
+					DeleteFunc: func(workqueue.RateLimitingInterface, event.DeleteEvent) {
+						defer GinkgoRecover()
+						Fail("Unexpected DeleteEvent")
+					},
+					GenericFunc: func(q2 workqueue.RateLimitingInterface, evt event.GenericEvent) {
+						defer GinkgoRecover()
+						time.Sleep(interval)
+					},
+				}, q)
+				Expect(err).NotTo(HaveOccurred())
+
+				// get channel blocked
+				ch <- evt
+				ch <- evt
+				ch <- evt
+
+				beforeEvent := time.Now()
+				ch <- evt
+				// validate event distribution get blocked.
+				distributeInterval := time.Now().Sub(beforeEvent)
+				Expect(distributeInterval >= interval).To(BeTrue())
+				close(done)
+			}, 15)
+			It("should get error if no source specified", func(done Done) {
+				q := workqueue.NewNamedRateLimitingQueue(workqueue.DefaultControllerRateLimiter(), "test")
+				instance := &source.ChannelSource{ /*no source specified*/ }
+				inject.DoStop(stop, instance)
+				err := instance.Start(eventhandler.Funcs{}, q)
+				Expect(err).To(Equal(fmt.Errorf("must specify ChannelSource.Source")))
+				close(done)
+			})
+			It("should get error if no stop channel injected", func(done Done) {
+				q := workqueue.NewNamedRateLimitingQueue(workqueue.DefaultControllerRateLimiter(), "test")
+				instance := &source.ChannelSource{Source: ch}
+				err := instance.Start(eventhandler.Funcs{}, q)
+				Expect(err).To(Equal(fmt.Errorf("must call InjectStop on ChannelSource before calling Start")))
+				close(done)
+			})
+
+		})
+		Context("for multi sources (handlers)", func() {
+			It("should provide GenericEvents for all handlers", func(done Done) {
+				ch := make(chan event.GenericEvent)
+				p := &corev1.Pod{
+					ObjectMeta: metav1.ObjectMeta{Name: "foo", Namespace: "bar"},
+				}
+				evt := event.GenericEvent{
+					Object: p,
+					Meta:   p,
+				}
+
+				var resEvent1, resEvent2 event.GenericEvent
+				c1 := make(chan struct{})
+				c2 := make(chan struct{})
+
+				q := workqueue.NewNamedRateLimitingQueue(workqueue.DefaultControllerRateLimiter(), "test")
+				instance := &source.ChannelSource{Source: ch}
+				inject.DoStop(stop, instance)
+				err := instance.Start(eventhandler.Funcs{
+					CreateFunc: func(workqueue.RateLimitingInterface, event.CreateEvent) {
+						defer GinkgoRecover()
+						Fail("Unexpected CreateEvent")
+					},
+					UpdateFunc: func(workqueue.RateLimitingInterface, event.UpdateEvent) {
+						defer GinkgoRecover()
+						Fail("Unexpected UpdateEvent")
+					},
+					DeleteFunc: func(workqueue.RateLimitingInterface, event.DeleteEvent) {
+						defer GinkgoRecover()
+						Fail("Unexpected DeleteEvent")
+					},
+					GenericFunc: func(q2 workqueue.RateLimitingInterface, evt event.GenericEvent) {
+						defer GinkgoRecover()
+						Expect(q2).To(Equal(q))
+						Expect(evt.Meta).To(Equal(p))
+						Expect(evt.Object).To(Equal(p))
+						resEvent1 = evt
+						close(c1)
+					},
+				}, q)
+				Expect(err).NotTo(HaveOccurred())
+
+				err = instance.Start(eventhandler.Funcs{
+					CreateFunc: func(workqueue.RateLimitingInterface, event.CreateEvent) {
+						defer GinkgoRecover()
+						Fail("Unexpected CreateEvent")
+					},
+					UpdateFunc: func(workqueue.RateLimitingInterface, event.UpdateEvent) {
+						defer GinkgoRecover()
+						Fail("Unexpected UpdateEvent")
+					},
+					DeleteFunc: func(workqueue.RateLimitingInterface, event.DeleteEvent) {
+						defer GinkgoRecover()
+						Fail("Unexpected DeleteEvent")
+					},
+					GenericFunc: func(q2 workqueue.RateLimitingInterface, evt event.GenericEvent) {
+						defer GinkgoRecover()
+						Expect(q2).To(Equal(q))
+						Expect(evt.Meta).To(Equal(p))
+						Expect(evt.Object).To(Equal(p))
+						resEvent2 = evt
+						close(c2)
+					},
+				}, q)
+				Expect(err).NotTo(HaveOccurred())
+
+				ch <- evt
+				<-c1
+				<-c2
+
+				// Validate the two handlers received same event
+				Expect(resEvent1).To(Equal(resEvent2))
+				close(done)
+			})
 		})
 	})
 })