Merge c23fe70 into ff8175b

api/v1/lib/extras/scheduler/controller/controller.go

@@ -21,9 +21,21 @@ type (
 		registrationTokens     <-chan struct{}
 		subscriptionTerminated func(error)
 		initSuppressRoles      []string
+		contextPerSubscription bool
 	}
 )
 
+// WithContextPerSubscription results in the creation of a sub-context that is passed to all event handlers
+// and is canceled when the associated subscription has termined (i.e. when the event loop exits and a re-
+// subscribe attempt is (possibly) attempted).
+func WithContextPerSubscription(b bool) Option {
+	return func(c *Config) Option {
+		old := c.contextPerSubscription
+		c.contextPerSubscription = b
+		return WithContextPerSubscription(old)
+	}
+}
+
 // WithInitiallySuppressedRoles sets the "suppressed_roles" field of the SUBSCRIBE call
 // that's issued to Mesos for each (re-)subscription attempt.
 func WithInitiallySuppressedRoles(r []string) Option {
@@ -140,6 +152,11 @@ func processSubscription(ctx context.Context, config Config, resp mesos.Response
 	if resp != nil {
 		defer resp.Close()
 	}
+	if config.contextPerSubscription {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithCancel(ctx)
+		defer cancel()
+	}
 	if err == nil {
 		err = eventLoop(ctx, config, resp)
 	}