<commit_before>// Copyright 2016 FullStory, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package solrmonitor

import (
	"errors"
	"reflect"
	"sync"
	"sync/atomic"
	"time"

	"github.com/samuel/go-zookeeper/zk"
)

// ugh... seems silly to have to define this
const maxInt = int(^uint(0) >> 1)

// A function that handles a zk event. Optionally returns a new event channel to automatically
// re-register.
type ZkEventHandler func(zk.Event) <-chan zk.Event

// An interface that handles a zk event. Optionally returns a new event channel to automatically
// re-register.
type ZkEventCallback interface {
	Handle(zk.Event) <-chan zk.Event
}

func (h *ZkEventHandler) Handle(event zk.Event) <-chan zk.Event {
	return (*h)(event)
}

// Monitors many zk.Event channels. Dispatches handling to other goroutines allowing handlers to
// run in parallel. A single ZkEventCallback, however, will always be invoked as if from the
// same goroutine. So a single callback can be used to watch multiple channels and it will receive
// events in a way that does not require it to synchronize with other possible invocations.
type ZkDispatcher struct {
	logger         zk.Logger            // where to debug log
	selectCases    []reflect.SelectCase // the list of event channels to watch
	selectHandlers []ZkEventCallback    // the list of handlers associated with each event channel
	runningProcs   int32                // for testing
	closed         int32                // interpreted as bool, defined as int32 for atomic ops
	closedChan     chan struct{}        // signals exit
	newHandlerChan chan newHandler      // sends new handlers to the event loop

	taskMu sync.Mutex
	tasks  map[ZkEventCallback]*fifoTaskQueue
}

type newHandler struct {
	watcher <-chan zk.Event
	handler ZkEventCallback
}

func NewZkDispatcher(logger zk.Logger) *ZkDispatcher {
	closedChan := make(chan struct{})
	newHandlerChan := make(chan newHandler, 1024)
	d := &ZkDispatcher{
		logger:         logger,
		closedChan:     closedChan,
		newHandlerChan: newHandlerChan,
		selectCases: []reflect.SelectCase{
			{Dir: reflect.SelectRecv, Chan: reflect.ValueOf(closedChan)},
			{Dir: reflect.SelectRecv, Chan: reflect.ValueOf(newHandlerChan)},
		},
		selectHandlers: []ZkEventCallback{nil, nil}, // first two correspond to close and new handler channels
		runningProcs:   1,
		tasks:          make(map[ZkEventCallback]*fifoTaskQueue),
	}
	go d.eventLoop()
	return d
}

func (d *ZkDispatcher) Close() {
	if !atomic.CompareAndSwapInt32(&d.closed, 0, 1) {
		panic("already closed")
	}

	close(d.closedChan)
}

var errClosed = errors.New("already closed")

// Watch a new ZK event using the given callback to handle the event.
func (d *ZkDispatcher) WatchEvent(watcher <-chan zk.Event, handler ZkEventCallback) error {
	if atomic.LoadInt32(&d.closed) != 0 {
		return errClosed
	}
	select {
	case d.newHandlerChan <- newHandler{watcher, handler}:
		return nil
	case <-d.closedChan:
		return errClosed
	case <-time.After(10 * time.Second):
		panic("channel is full")
	}
	return nil
}

// Watch a new ZK event using the given function to handle the event.
func (d *ZkDispatcher) Watch(watcher <-chan zk.Event, handler ZkEventHandler) error {
	return d.WatchEvent(watcher, &handler)
}

func (d *ZkDispatcher) eventLoop() {
	defer atomic.AddInt32(&d.runningProcs, -1)

	for {
		// first try to drain any new handlers, bailing if dispatcher is closed
		for done := false; !done; {
			select {
			case nh := <-d.newHandlerChan:
				d.selectHandlers = append(d.selectHandlers, nh.handler)
				d.selectCases = append(d.selectCases, newCase(nh.watcher))
			case <-d.closedChan:
				d.logger.Printf("zkdispatcher: exiting via close channel")
				return
			default:
				done = true // nothing else waiting
			}
		}

		// then try to select an event from one of the watched channels<commit_msg>Use the close channel to signal closure; remove close atomic (#11)<commit_after>// Copyright 2016 FullStory, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package solrmonitor

import (
	"errors"
	"reflect"
	"sync"
	"sync/atomic"
	"time"

	"github.com/samuel/go-zookeeper/zk"
)

// ugh... seems silly to have to define this
const maxInt = int(^uint(0) >> 1)

// A function that handles a zk event. Optionally returns a new event channel to automatically
// re-register.
type ZkEventHandler func(zk.Event) <-chan zk.Event

// An interface that handles a zk event. Optionally returns a new event channel to automatically
// re-register.
type ZkEventCallback interface {
	Handle(zk.Event) <-chan zk.Event
}

func (h *ZkEventHandler) Handle(event zk.Event) <-chan zk.Event {
	return (*h)(event)
}

// Monitors many zk.Event channels. Dispatches handling to other goroutines allowing handlers to
// run in parallel. A single ZkEventCallback, however, will always be invoked as if from the
// same goroutine. So a single callback can be used to watch multiple channels and it will receive
// events in a way that does not require it to synchronize with other possible invocations.
type ZkDispatcher struct {
	logger         zk.Logger            // where to debug log
	selectCases    []reflect.SelectCase // the list of event channels to watch
	selectHandlers []ZkEventCallback    // the list of handlers associated with each event channel
	runningProcs   int32                // for testing
	closedChan     chan struct{}        // signals exit
	newHandlerChan chan newHandler      // sends new handlers to the event loop

	taskMu sync.Mutex
	tasks  map[ZkEventCallback]*fifoTaskQueue
}

type newHandler struct {
	watcher <-chan zk.Event
	handler ZkEventCallback
}

func NewZkDispatcher(logger zk.Logger) *ZkDispatcher {
	closedChan := make(chan struct{})
	newHandlerChan := make(chan newHandler, 1024)
	d := &ZkDispatcher{
		logger:         logger,
		closedChan:     closedChan,
		newHandlerChan: newHandlerChan,
		selectCases: []reflect.SelectCase{
			{Dir: reflect.SelectRecv, Chan: reflect.ValueOf(closedChan)},
			{Dir: reflect.SelectRecv, Chan: reflect.ValueOf(newHandlerChan)},
		},
		selectHandlers: []ZkEventCallback{nil, nil}, // first two correspond to close and new handler channels
		runningProcs:   1,
		tasks:          make(map[ZkEventCallback]*fifoTaskQueue),
	}
	go d.eventLoop()
	return d
}

func (d *ZkDispatcher) Close() {
	close(d.closedChan)
}

var errClosed = errors.New("already closed")

// Watch a new ZK event using the given callback to handle the event.
func (d *ZkDispatcher) WatchEvent(watcher <-chan zk.Event, handler ZkEventCallback) error {
	// Check for closure first
	select {
	case <-d.closedChan:
		return errClosed
	default:
		break
	}

	select {
	case <-d.closedChan:
		return errClosed
	case d.newHandlerChan <- newHandler{watcher, handler}:
		return nil
	case <-time.After(10 * time.Second):
		// Escape hatch to avoid deadlock
		panic("channel is full")
	}
	return nil
}

// Watch a new ZK event using the given function to handle the event.
func (d *ZkDispatcher) Watch(watcher <-chan zk.Event, handler ZkEventHandler) error {
	return d.WatchEvent(watcher, &handler)
}

func (d *ZkDispatcher) eventLoop() {
	defer atomic.AddInt32(&d.runningProcs, -1)

	for {
		// first try to drain any new handlers, bailing if dispatcher is closed
		for done := false; !done; {
			select {
			case nh := <-d.newHandlerChan:
				d.selectHandlers = append(d.selectHandlers, nh.handler)
				d.selectCases = append(d.selectCases, newCase(nh.watcher))
			case <-d.closedChan:
				d.logger.Printf("zkdispatcher: exiting via close channel")
				return
			default:
				done = true // nothing else waiting
			}
		}

		// then try to select an event from one of the watched channels
{'fraction_non_alphanumeric': 0.09550756438387804, 'fraction_numerical': 0.006582746275551449, 'mean_word_length': 3.567510548523207, 'pattern_counts': {'":': 0, '<': 26, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5366898', 'n_tokens_mistral': 2809, 'n_tokens_neox': 2566, 'n_words': 965}