<commit_before>/* Licensed to the Apache Software Foundation (ASF) under one or more
contributor license agreements.  See the NOTICE file distributed with
this work for additional information regarding copyright ownership.
The ASF licenses this file to You 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 go_kafka_client

import (
	"fmt"
	"math"
	"sync"
	"time"
)

type consumerFetcherManager struct {
	config                         *ConsumerConfig
	numStreams                     int
	closeFinished                  chan bool
	updateLock                     sync.RWMutex
	partitionMap                   map[TopicAndPartition]*partitionTopicInfo
	fetcherRoutineMap              map[int]*consumerFetcherRoutine
	shuttingDown                   bool
	updateInProgress               bool
	updatedCond                    *sync.Cond
	disconnectChannelsForPartition chan TopicAndPartition

	metrics *ConsumerMetrics
	client  LowLevelClient
}

func (m *consumerFetcherManager) String() string {
	return fmt.Sprintf("%s-manager", m.config.Consumerid)
}

func newConsumerFetcherManager(config *ConsumerConfig, disconnectChannelsForPartition chan TopicAndPartition, metrics *ConsumerMetrics) *consumerFetcherManager {
	manager := &consumerFetcherManager{
		config:                         config,
		closeFinished:                  make(chan bool),
		partitionMap:                   make(map[TopicAndPartition]*partitionTopicInfo),
		fetcherRoutineMap:              make(map[int]*consumerFetcherRoutine),
		disconnectChannelsForPartition: disconnectChannelsForPartition,
		client:  config.LowLevelClient,
		metrics: metrics,
	}
	manager.updatedCond = sync.NewCond(manager.updateLock.RLocker())

	return manager
}

func (m *consumerFetcherManager) startConnections(topicInfos []*partitionTopicInfo, numStreams int) {
	Debug(m, "Fetcher Manager started")
	Debugf(m, "TopicInfos = %s", topicInfos)
	m.numStreams = numStreams

	partitionInfos := make(map[TopicAndPartition]*partitionTopicInfo)

	m.updateInProgress = true
	inWriteLock(&m.updateLock, func() {
		Debug(m, "Updating fetcher configuration")
		newPartitionMap := make(map[TopicAndPartition]*partitionTopicInfo)
		for _, providedInfo := range topicInfos {
			topicAndPartition := TopicAndPartition{providedInfo.Topic, providedInfo.Partition}
			if currentInfo, alreadyFetching := m.partitionMap[topicAndPartition]; alreadyFetching {
				newPartitionMap[topicAndPartition] = currentInfo
				continue
			} else {
				newPartitionMap[topicAndPartition] = providedInfo
				partitionInfos[topicAndPartition] = providedInfo
			}
		}

		Debugf(m, "Got new list of partitions to process %v", newPartitionMap)
		Debugf(m, "All partitions map: %v", m.partitionMap)
		//receive obsolete partitions map
		for k := range newPartitionMap {
			delete(m.partitionMap, k)
		}

		//removing obsolete partitions and tearing down associated jobs
		topicPartitionsToRemove := make([]TopicAndPartition, 0)
		for tp := range m.partitionMap {
			topicPartitionsToRemove = append(topicPartitionsToRemove, tp)
		}
		Debugf(m, "There are obsolete partitions %v", topicPartitionsToRemove)

		//removing unnecessary partition-fetchRoutine bindings
		for _, fetcher := range m.fetcherRoutineMap {
			Debugf(m, "Fetcher %s parition map before obsolete partitions removal", fetcher, fetcher.partitionMap)
			fetcher.removePartitions(topicPartitionsToRemove)
			Debugf(m, "Fetcher %s parition map after obsolete partitions removal", fetcher, fetcher.partitionMap)
		}
		for tp := range m.partitionMap {
			m.disconnectChannelsForPartition <- tp
			delete(m.partitionMap, tp)
		}
		m.shutdownIdleFetchers()

		//updating partitions map with requested partitions
		for k, v := range newPartitionMap {
			m.partitionMap[k] = v
		}
		m.addFetcherForPartitions(partitionInfos)

		m.updateInProgress = false
		Debugf(m, "Applied new partition map %v", m.partitionMap)
	})

	Debug(m, "Notifying all waiters about completed update")
	m.updatedCond.Broadcast()
}

func (m *consumerFetcherManager) addFetcherForPartitions(partitionInfos map[TopicAndPartition]*partitionTopicInfo) {
	Infof(m, "Adding fetcher for partitions %v", partitionInfos)
	partitionsPerFetcher := make(map[int]map[TopicAndPartition]*partitionTopicInfo)
	for topicAndPartition, info := range partitionInfos {
		fetcherId := m.getFetcherId(topicAndPartition.Topic, topicAndPartition.Partition)
		if partitionsPerFetcher[fetcherId] == nil {
			partitionsPerFetcher[fetcherId] = make(map[TopicAndPartition]*partitionTopicInfo)
		}
		partitionsPerFetcher[fetcherId][topicAndPartition] = info
	}

	Debugf(m, "partitionsPerFetcher: %v", partitionsPerFetcher)
	for fetcherId, partitionInfos := range partitionsPerFetcher {
		if m.fetcherRoutineMap[fetcherId] == nil {
			Debugf(m, "Starting new fetcher")
			fetcherRoutine := newConsumerFetcher(m,
				fmt.Sprintf("ConsumerFetcherRoutine-%s-%d", m.config.Consumerid, fetcherId))
			m.fetcherRoutineMap[fetcherId] = fetcherRoutine
			go fetcherRoutine.start()
		}

		m.fetcherRoutineMap[fetcherId].addPartitions(partitionInfos)
	}
}

func (m *consumerFetcherManager) getFetcherId(topic string, partitionId int32) int {
	return int(math.Abs(float64(31*hash(topic)+partitionId))) % int(m.numStreams)
}

func (m *consumerFetcherManager) shutdownIdleFetchers() {
	Debug(m, "Shutting down idle fetchers")
	for key, fetcher := range m.fetcherRoutineMap {
		if len(fetcher.partitionMap) <= 0 {
			Debugf(m, "There is idle fetcher: %s", fetcher)
			<-fetcher.close()
			delete(m.fetcherRoutineMap, key)
		}
	}
	Debug(m, "Closed idle fetchers")
}

func (m *consumerFetcherManager) closeAllFetchers() {
	Info(m, "Closing fetchers")
	inWriteLock(&m.updateLock, func() {
		Debug(m, "Stopping all message buffers and removing paritions")
		for k, v := range m.partitionMap {
			v.Buffer.stop()
			delete(m.partitionMap, k)
		}

		Debugf(m, "Trying to close %d fetchers", len(m.fetcherRoutineMap))
		for key, fetcher := range m.fetcherRoutineMap {
			Debugf(m, "Closing %s", fetcher)
			<-fetcher.close()
			Debugf(m, "Closed %s", fetcher)
			delete(m.fetcherRoutineMap, key)
		}
	})
}

func (m *consumerFetcherManager) close() <-chan bool {
	Info(m, "Closing manager")
	go func() {
		Info(m, "Setting shutdown flag")
		m.shuttingDown = true
		m.closeAllFetchers()
		m.updatedCond.Broadcast()
		m.partitionMap = nil
		m.closeFinished <- true
		Info(m, "Successfully closed all fetcher manager routines")
	}()

	return m.closeFinished
}

type consumerFetcherRoutine struct {
	manager       *consumerFetcherManager
	name          string
	partitionMap  map[TopicAndPartition]*partitionTopicInfo
	lock          sync.RWMutex
	closeFinished chan bool
	fetchStopper  chan bool
	askNext       chan TopicAndPartition
}

func (f *consumerFetcherRoutine) String() string {
	return f.name
}

func newConsumerFetcher(m *consumerFetcherManager, name string) *consumerFetcherRoutine {
	return &consumerFetcherRoutine{
		manager:       m,
		name:          name,
		partitionMap:  make(map[TopicAndPartition]*partitionTopicInfo),
		closeFinished: make(chan bool),
		fetchStopper:  make(chan bool),
		askNext:       make(chan TopicAndPartition, m.config.AskNextChannelSize),
	}
}

func (f *consumerFetcherRoutine) start() {
	Info(f, "Fetcher started")
	for {
		if Logger.IsAllowed(TraceLevel) {
			Trace(f, "Waiting for asknext or die")
		}
		ts := time.Now()
		select {
		case nextTopicPartition := <-f.askNext:
			{
				timestamp := time.Now().UnixNano() / int64(time.Millisecond)
				f.manager.metrics.fetchersIdle().Update(time.Since(ts))
				Debugf(f, "Received asknext for %s", &nextTopicPartition)
				inReadLock(&f.lock, func() {
					if !f.manager.shuttingDown {
						Debugf(f, "Partition map: %v", f.partitionMap)
						if _, exists := f.partitionMap[nextTopicPartition]; !exists {
							Warnf(f, "Message buffer for partition %s has been terminated. Aborting processing task...", nextTopicPartition)
							return
						}
						offset := f.partitionMap[nextTopicPartition].FetchedOffset

						var messages []*Message
						var err error
						f.manager.metrics.fetchDuration().Time(func() {
							messages, err = f.manager.client.Fetch(nextTopicPartition.Topic, nextTopicPartition.Partition, offset)
						})

						if err != nil {
							if offset > -1 { // Negative offsets are obviously out of range but don't spam the logs...
								if f.manager.client.IsOffsetOutOfRange(err) {
									Warnf(f, "Current offset %d for topic %s and partition %s is out of range.", offset, nextTopicPartition.Topic, nextTopicPartition.Partition)
									f.handleOffsetOutOfRange(&nextTopicPartition)
								} else {
									Warnf(f, "Got a fetch error: %s", err)
									//TODO new backoff type?
									time.Sleep(1 * time.Second)
								}
							}
						}

						if f.manager.config.Debug {
							for _, message := range messages {
								message.DecodedKey = append([]int64{timestamp}, message.DecodedKey.([]int64)...)
							}
						}

						f.processPartitionData(nextTopicPartition, messages)
					}
				})
			}
		case <-f.fetchStopper:
			{
				Info(f, "Stopped fetcher")
				return
			}
		}
	}
}

func (f *consumerFetcherRoutine) addPartitions(partitionTopicInfos map[TopicAndPartition]*partitionTopicInfo) {
	Debugf(f, "Adding partitions: %v", partitionTopicInfos)
	newPartitions := make(map[TopicAndPartition]chan TopicAndPartition)
	inWriteLock(&f.lock, func() {
		for topicAndPartition, info := range partitionTopicInfos {
			if _, contains := f.partitionMap[topicAndPartition]; !contains {
				f.partitionMap[topicAndPartition] = info
				validOffset := info.FetchedOffset + 1
				if isOffsetInvalid(info.FetchedOffset) {
					f.handleOffsetOutOfRange(&topicAndPartition)
				} else {
					f.partitionMap[topicAndPartition].FetchedOffset = validOffset
				}
				f.partitionMap[topicAndPartition].Buffer.start(f.askNext)
				newPartitions[topicAndPartition] = f.askNext
				Debugf(f, "Owner of %s", topicAndPartition)
			}
		}
	})

	for topicAndPartition, askNext := range newPartitions {
		Debugf(f, "Sending ask next to %s for %s", f, topicAndPartition)
	Loop:
		for {
			timeout := time.NewTimer(1 * time.Second)
			select {
			case askNext <- topicAndPartition:
				timeout.Stop()
				break Loop
			case <-timeout.C:
				{
					if f.manager.shuttingDown {
						return
					}
				}
			}
		}
		Debugf(f, "Sent ask next to %s for %s", f, topicAndPartition)
	}
}

func (f *consumerFetcherRoutine) processPartitionData(topicAndPartition TopicAndPartition, messages []*Message) {
	if Logger.IsAllowed(TraceLevel) {
		Trace(f, "Trying to acquire lock for partition processing")
		Tracef(f, "Processing partition data for %s", topicAndPartition)
	}
	if len(messages) > 0 {
		f.partitionMap[topicAndPartition].FetchedOffset = messages[len(messages)-1].Offset + 1
	}
	go f.partitionMap[topicAndPartition].Buffer.addBatch(messages)
	if Logger.IsAllowed(TraceLevel) {
		Tracef(f, "Sent partition data to %s", topicAndPartition)
	}
}

func (f *consumerFetcherRoutine) handleOffsetOutOfRange(topicAndPartition *TopicAndPartition) {
	newOffset, err := f.manager.client.GetAvailableOffset(topicAndPartition.Topic, topicAndPartition.Partition, f.manager.config.AutoOffsetReset)
	if err != nil {
		Errorf(f, "Cannot get available offset for %s. Reason: %s", topicAndPartition, err)
		return
	}

	// Do not use a lock here just because it's faster and it will be checked afterwards if we should still fetch that TopicPartition
	// This just guarantees we dont get a nil pointer dereference here
	if topicInfo, exists := f.partitionMap[*topicAndPartition]; exists {
		topicInfo.FetchedOffset = newOffset
		f.partitionMap[*topicAndPartition].FetchedOffset = newOffset
	}
}

func (f *consumerFetcherRoutine) removeAllPartitions() {
	partitions := make([]TopicAndPartition, 0)
	for topicPartition, _ := range f.partitionMap {
		partitions = append(partitions, topicPartition)
	}
	f.removePartitions(partitions)
}

func (f *consumerFetcherRoutine) removePartitions(partitions []TopicAndPartition) {
	Debug(f, "Remove partitions")
	inWriteLock(&f.lock, func() {
		for _, topicAndPartition := range partitions {
			delete(f.partitionMap, topicAndPartition)
		}
	})
}

func (f *consumerFetcherRoutine) close() <-chan bool {
	Info(f, "Closing fetcher")
	go func() {
		f.fetchStopper <- true
		f.removeAllPartitions()
		Debug(f, "Sending close finished")
		f.closeFinished <- true
		Debug(f, "Sent close finished")
	}()
	return f.closeFinished
}
<commit_msg>Added topic-partition logging on fetch errors<commit_after>/* Licensed to the Apache Software Foundation (ASF) under one or more
contributor license agreements.  See the NOTICE file distributed with
this work for additional information regarding copyright ownership.
The ASF licenses this file to You 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 go_kafka_client

import (
	"fmt"
	"math"
	"sync"
	"time"
)

type consumerFetcherManager struct {
	config                         *ConsumerConfig
	numStreams                     int
	closeFinished                  chan bool
	updateLock                     sync.RWMutex
	partitionMap                   map[TopicAndPartition]*partitionTopicInfo
	fetcherRoutineMap              map[int]*consumerFetcherRoutine
	shuttingDown                   bool
	updateInProgress               bool
	updatedCond                    *sync.Cond
	disconnectChannelsForPartition chan TopicAndPartition

	metrics *ConsumerMetrics
	client  LowLevelClient
}

func (m *consumerFetcherManager) String() string {
	return fmt.Sprintf("%s-manager", m.config.Consumerid)
}

func newConsumerFetcherManager(config *ConsumerConfig, disconnectChannelsForPartition chan TopicAndPartition, metrics *ConsumerMetrics) *consumerFetcherManager {
	manager := &consumerFetcherManager{
		config:                         config,
		closeFinished:                  make(chan bool),
		partitionMap:                   make(map[TopicAndPartition]*partitionTopicInfo),
		fetcherRoutineMap:              make(map[int]*consumerFetcherRoutine),
		disconnectChannelsForPartition: disconnectChannelsForPartition,
		client:  config.LowLevelClient,
		metrics: metrics,
	}
	manager.updatedCond = sync.NewCond(manager.updateLock.RLocker())

	return manager
}

func (m *consumerFetcherManager) startConnections(topicInfos []*partitionTopicInfo, numStreams int) {
	Debug(m, "Fetcher Manager started")
	Debugf(m, "TopicInfos = %s", topicInfos)
	m.numStreams = numStreams

	partitionInfos := make(map[TopicAndPartition]*partitionTopicInfo)

	m.updateInProgress = true
	inWriteLock(&m.updateLock, func() {
		Debug(m, "Updating fetcher configuration")
		newPartitionMap := make(map[TopicAndPartition]*partitionTopicInfo)
		for _, providedInfo := range topicInfos {
			topicAndPartition := TopicAndPartition{providedInfo.Topic, providedInfo.Partition}
			if currentInfo, alreadyFetching := m.partitionMap[topicAndPartition]; alreadyFetching {
				newPartitionMap[topicAndPartition] = currentInfo
				continue
			} else {
				newPartitionMap[topicAndPartition] = providedInfo
				partitionInfos[topicAndPartition] = providedInfo
			}
		}

		Debugf(m, "Got new list of partitions to process %v", newPartitionMap)
		Debugf(m, "All partitions map: %v", m.partitionMap)
		//receive obsolete partitions map
		for k := range newPartitionMap {
			delete(m.partitionMap, k)
		}

		//removing obsolete partitions and tearing down associated jobs
		topicPartitionsToRemove := make([]TopicAndPartition, 0)
		for tp := range m.partitionMap {
			topicPartitionsToRemove = append(topicPartitionsToRemove, tp)
		}
		Debugf(m, "There are obsolete partitions %v", topicPartitionsToRemove)

		//removing unnecessary partition-fetchRoutine bindings
		for _, fetcher := range m.fetcherRoutineMap {
			Debugf(m, "Fetcher %s parition map before obsolete partitions removal", fetcher, fetcher.partitionMap)
			fetcher.removePartitions(topicPartitionsToRemove)
			Debugf(m, "Fetcher %s parition map after obsolete partitions removal", fetcher, fetcher.partitionMap)
		}
		for tp := range m.partitionMap {
			m.disconnectChannelsForPartition <- tp
			delete(m.partitionMap, tp)
		}
		m.shutdownIdleFetchers()

		//updating partitions map with requested partitions
		for k, v := range newPartitionMap {
			m.partitionMap[k] = v
		}
		m.addFetcherForPartitions(partitionInfos)

		m.updateInProgress = false
		Debugf(m, "Applied new partition map %v", m.partitionMap)
	})

	Debug(m, "Notifying all waiters about completed update")
	m.updatedCond.Broadcast()
}

func (m *consumerFetcherManager) addFetcherForPartitions(partitionInfos map[TopicAndPartition]*partitionTopicInfo) {
	Infof(m, "Adding fetcher for partitions %v", partitionInfos)
	partitionsPerFetcher := make(map[int]map[TopicAndPartition]*partitionTopicInfo)
	for topicAndPartition, info := range partitionInfos {
		fetcherId := m.getFetcherId(topicAndPartition.Topic, topicAndPartition.Partition)
		if partitionsPerFetcher[fetcherId] == nil {
			partitionsPerFetcher[fetcherId] = make(map[TopicAndPartition]*partitionTopicInfo)
		}
		partitionsPerFetcher[fetcherId][topicAndPartition] = info
	}

	Debugf(m, "partitionsPerFetcher: %v", partitionsPerFetcher)
	for fetcherId, partitionInfos := range partitionsPerFetcher {
		if m.fetcherRoutineMap[fetcherId] == nil {
			Debugf(m, "Starting new fetcher")
			fetcherRoutine := newConsumerFetcher(m,
				fmt.Sprintf("ConsumerFetcherRoutine-%s-%d", m.config.Consumerid, fetcherId))
			m.fetcherRoutineMap[fetcherId] = fetcherRoutine
			go fetcherRoutine.start()
		}

		m.fetcherRoutineMap[fetcherId].addPartitions(partitionInfos)
	}
}

func (m *consumerFetcherManager) getFetcherId(topic string, partitionId int32) int {
	return int(math.Abs(float64(31*hash(topic)+partitionId))) % int(m.numStreams)
}

func (m *consumerFetcherManager) shutdownIdleFetchers() {
	Debug(m, "Shutting down idle fetchers")
	for key, fetcher := range m.fetcherRoutineMap {
		if len(fetcher.partitionMap) <= 0 {
			Debugf(m, "There is idle fetcher: %s", fetcher)
			<-fetcher.close()
			delete(m.fetcherRoutineMap, key)
		}
	}
	Debug(m, "Closed idle fetchers")
}

func (m *consumerFetcherManager) closeAllFetchers() {
	Info(m, "Closing fetchers")
	inWriteLock(&m.updateLock, func() {
		Debug(m, "Stopping all message buffers and removing paritions")
		for k, v := range m.partitionMap {
			v.Buffer.stop()
			delete(m.partitionMap, k)
		}

		Debugf(m, "Trying to close %d fetchers", len(m.fetcherRoutineMap))
		for key, fetcher := range m.fetcherRoutineMap {
			Debugf(m, "Closing %s", fetcher)
			<-fetcher.close()
			Debugf(m, "Closed %s", fetcher)
			delete(m.fetcherRoutineMap, key)
		}
	})
}

func (m *consumerFetcherManager) close() <-chan bool {
	Info(m, "Closing manager")
	go func() {
		Info(m, "Setting shutdown flag")
		m.shuttingDown = true
		m.closeAllFetchers()
		m.updatedCond.Broadcast()
		m.partitionMap = nil
		m.closeFinished <- true
		Info(m, "Successfully closed all fetcher manager routines")
	}()

	return m.closeFinished
}

type consumerFetcherRoutine struct {
	manager       *consumerFetcherManager
	name          string
	partitionMap  map[TopicAndPartition]*partitionTopicInfo
	lock          sync.RWMutex
	closeFinished chan bool
	fetchStopper  chan bool
	askNext       chan TopicAndPartition
}

func (f *consumerFetcherRoutine) String() string {
	return f.name
}

func newConsumerFetcher(m *consumerFetcherManager, name string) *consumerFetcherRoutine {
	return &consumerFetcherRoutine{
		manager:       m,
		name:          name,
		partitionMap:  make(map[TopicAndPartition]*partitionTopicInfo),
		closeFinished: make(chan bool),
		fetchStopper:  make(chan bool),
		askNext:       make(chan TopicAndPartition, m.config.AskNextChannelSize),
	}
}

func (f *consumerFetcherRoutine) start() {
	Info(f, "Fetcher started")
	for {
		if Logger.IsAllowed(TraceLevel) {
			Trace(f, "Waiting for asknext or die")
		}
		ts := time.Now()
		select {
		case nextTopicPartition := <-f.askNext:
			{
				timestamp := time.Now().UnixNano() / int64(time.Millisecond)
				f.manager.metrics.fetchersIdle().Update(time.Since(ts))
				Debugf(f, "Received asknext for %s", &nextTopicPartition)
				inReadLock(&f.lock, func() {
					if !f.manager.shuttingDown {
						Debugf(f, "Partition map: %v", f.partitionMap)
						if _, exists := f.partitionMap[nextTopicPartition]; !exists {
							Warnf(f, "Message buffer for partition %s has been terminated. Aborting processing task...", nextTopicPartition)
							return
						}
						offset := f.partitionMap[nextTopicPartition].FetchedOffset

						var messages []*Message
						var err error
						f.manager.metrics.fetchDuration().Time(func() {
							messages, err = f.manager.client.Fetch(nextTopicPartition.Topic, nextTopicPartition.Partition, offset)
						})

						if err != nil {
							if offset > -1 { // Negative offsets are obviously out of range but don't spam the logs...
								if f.manager.client.IsOffsetOutOfRange(err) {
									Warnf(f, "Current offset %d for topic %s and partition %s is out of range.", offset, nextTopicPartition.Topic, nextTopicPartition.Partition)
									f.handleOffsetOutOfRange(&nextTopicPartition)
								} else {
									Warnf(f, "Got a fetch error for topic %s, partition %d: %s", nextTopicPartition.Topic, nextTopicPartition.Partition, err)
									//TODO new backoff type?
									time.Sleep(1 * time.Second)
								}
							}
						}

						if f.manager.config.Debug {
							for _, message := range messages {
								message.DecodedKey = append([]int64{timestamp}, message.DecodedKey.([]int64)...)
							}
						}

						f.processPartitionData(nextTopicPartition, messages)
					}
				})
			}
		case <-f.fetchStopper:
			{
				Info(f, "Stopped fetcher")
				return
			}
		}
	}
}

func (f *consumerFetcherRoutine) addPartitions(partitionTopicInfos map[TopicAndPartition]*partitionTopicInfo) {
	Debugf(f, "Adding partitions: %v", partitionTopicInfos)
	newPartitions := make(map[TopicAndPartition]chan TopicAndPartition)
	inWriteLock(&f.lock, func() {
		for topicAndPartition, info := range partitionTopicInfos {
			if _, contains := f.partitionMap[topicAndPartition]; !contains {
				f.partitionMap[topicAndPartition] = info
				validOffset := info.FetchedOffset + 1
				if isOffsetInvalid(info.FetchedOffset) {
					f.handleOffsetOutOfRange(&topicAndPartition)
				} else {
					f.partitionMap[topicAndPartition].FetchedOffset = validOffset
				}
				f.partitionMap[topicAndPartition].Buffer.start(f.askNext)
				newPartitions[topicAndPartition] = f.askNext
				Debugf(f, "Owner of %s", topicAndPartition)
			}
		}
	})

	for topicAndPartition, askNext := range newPartitions {
		Debugf(f, "Sending ask next to %s for %s", f, topicAndPartition)
	Loop:
		for {
			timeout := time.NewTimer(1 * time.Second)
			select {
			case askNext <- topicAndPartition:
				timeout.Stop()
				break Loop
			case <-timeout.C:
				{
					if f.manager.shuttingDown {
						return
					}
				}
			}
		}
		Debugf(f, "Sent ask next to %s for %s", f, topicAndPartition)
	}
}

func (f *consumerFetcherRoutine) processPartitionData(topicAndPartition TopicAndPartition, messages []*Message) {
	if Logger.IsAllowed(TraceLevel) {
		Trace(f, "Trying to acquire lock for partition processing")
		Tracef(f, "Processing partition data for %s", topicAndPartition)
	}
	if len(messages) > 0 {
		f.partitionMap[topicAndPartition].FetchedOffset = messages[len(messages)-1].Offset + 1
	}
	go f.partitionMap[topicAndPartition].Buffer.addBatch(messages)
	if Logger.IsAllowed(TraceLevel) {
		Tracef(f, "Sent partition data to %s", topicAndPartition)
	}
}

func (f *consumerFetcherRoutine) handleOffsetOutOfRange(topicAndPartition *TopicAndPartition) {
	newOffset, err := f.manager.client.GetAvailableOffset(topicAndPartition.Topic, topicAndPartition.Partition, f.manager.config.AutoOffsetReset)
	if err != nil {
		Errorf(f, "Cannot get available offset for %s. Reason: %s", topicAndPartition, err)
		return
	}

	// Do not use a lock here just because it's faster and it will be checked afterwards if we should still fetch that TopicPartition
	// This just guarantees we dont get a nil pointer dereference here
	if topicInfo, exists := f.partitionMap[*topicAndPartition]; exists {
		topicInfo.FetchedOffset = newOffset
		f.partitionMap[*topicAndPartition].FetchedOffset = newOffset
	}
}

func (f *consumerFetcherRoutine) removeAllPartitions() {
	partitions := make([]TopicAndPartition, 0)
	for topicPartition, _ := range f.partitionMap {
		partitions = append(partitions, topicPartition)
	}
	f.removePartitions(partitions)
}

func (f *consumerFetcherRoutine) removePartitions(partitions []TopicAndPartition) {
	Debug(f, "Remove partitions")
	inWriteLock(&f.lock, func() {
		for _, topicAndPartition := range partitions {
			delete(f.partitionMap, topicAndPartition)
		}
	})
}

func (f *consumerFetcherRoutine) close() <-chan bool {
	Info(f, "Closing fetcher")
	go func() {
		f.fetchStopper <- true
		f.removeAllPartitions()
		Debug(f, "Sending close finished")
		f.closeFinished <- true
		Debug(f, "Sent close finished")
	}()
	return f.closeFinished
}
<|endoftext|>


<commit_before>package dispatcher

import (
	"net"
	"sync"
	"time"
)

const blockDuration = 15 * time.Minute

// GateKeeper controls whether peers with a given address may connect with us.
type GateKeeper struct {
	addresses sync.Map
	done      chan bool
}

// NewGateKeeper creates a new GateKeeper.
func NewGateKeeper() *GateKeeper {
	r := &GateKeeper{done: make(chan bool)}
	go r.prune()
	return r
}

// BlockAddress adds the specified address to the incoming blocked list.
func (r *GateKeeper) BlockAddress(addr net.Addr) {
	host, _, err := net.SplitHostPort(addr.String())
	if err != nil {
		return
	}
	r.BlockAddressString(host)
}

// BlockAddressString adds the specified address to the incoming blocked list.
func (r *GateKeeper) BlockAddressString(addr string) {
	r.addresses.Store(addr, time.Now().Add(blockDuration))
}

// IsAddressBlocked returns true if the address is blocked.
func (r *GateKeeper) IsAddressBlocked(addr net.Addr) bool {
	host, _, err := net.SplitHostPort(addr.String())
	if err != nil {
		return true
	}
	return r.IsAddressStringBlocked(host)
}

// IsAddressStringBlocked returns true if the address is blocked.
func (r *GateKeeper) IsAddressStringBlocked(addr string) bool {
	expires, ok := r.addresses.Load(addr)
	return ok && !expires.(time.Time).Before(time.Now())
}

func (r *GateKeeper) prune() {
	for {
		select {
		case <-time.After(blockDuration):
			r.addresses.Range(func(addr interface{}, expires interface{}) bool {
				if expires.(time.Time).Before(time.Now()) {
					r.addresses.Delete(addr)
				}
				return true
			})
		case <-r.done:
			return
		}
	}
}

// Close shuts this GateKeeper down.
func (r *GateKeeper) Close() {
	close(r.done)
}
<commit_msg>Reduce block duration<commit_after>package dispatcher

import (
	"net"
	"sync"
	"time"
)

const blockDuration = 5 * time.Minute

// GateKeeper controls whether peers with a given address may connect with us.
type GateKeeper struct {
	addresses sync.Map
	done      chan bool
}

// NewGateKeeper creates a new GateKeeper.
func NewGateKeeper() *GateKeeper {
	r := &GateKeeper{done: make(chan bool)}
	go r.prune()
	return r
}

// BlockAddress adds the specified address to the incoming blocked list.
func (r *GateKeeper) BlockAddress(addr net.Addr) {
	host, _, err := net.SplitHostPort(addr.String())
	if err != nil {
		return
	}
	r.BlockAddressString(host)
}

// BlockAddressString adds the specified address to the incoming blocked list.
func (r *GateKeeper) BlockAddressString(addr string) {
	r.addresses.Store(addr, time.Now().Add(blockDuration))
}

// IsAddressBlocked returns true if the address is blocked.
func (r *GateKeeper) IsAddressBlocked(addr net.Addr) bool {
	host, _, err := net.SplitHostPort(addr.String())
	if err != nil {
		return true
	}
	return r.IsAddressStringBlocked(host)
}

// IsAddressStringBlocked returns true if the address is blocked.
func (r *GateKeeper) IsAddressStringBlocked(addr string) bool {
	expires, ok := r.addresses.Load(addr)
	return ok && !expires.(time.Time).Before(time.Now())
}

func (r *GateKeeper) prune() {
	for {
		select {
		case <-time.After(blockDuration):
			r.addresses.Range(func(addr interface{}, expires interface{}) bool {
				if expires.(time.Time).Before(time.Now()) {
					r.addresses.Delete(addr)
				}
				return true
			})
		case <-r.done:
			return
		}
	}
}

// Close shuts this GateKeeper down.
func (r *GateKeeper) Close() {
	close(r.done)
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2011 Matt Jibson <matt.jibson@gmail.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

// Package fft provides forward and inverse fast Fourier transform functions.
package fft

import (
	"math"
	"runtime"
	"sync"

	"github.com/mjibson/go-dsp/dsputils"
)

var (
	radix2Lock    sync.RWMutex
	radix2Factors = map[int][]complex128{
		4: {complex(1, 0), complex(0, -1), complex(-1, 0), complex(0, 1)},
	}
)

var (
	bluesteinLock       sync.RWMutex
	bluesteinFactors    = map[int][]complex128{}
	bluesteinInvFactors = map[int][]complex128{}
)

// EnsureRadix2Factors ensures that all radix 2 factors are computed for inputs
// of length input_len. This is used to precompute needed factors for known
// sizes. Generally should only be used for benchmarks.
func EnsureRadix2Factors(input_len int) {
	getRadix2Factors(input_len)
}

func getRadix2Factors(input_len int) []complex128 {
	radix2Lock.RLock()

	if hasRadix2Factors(input_len) {
		defer radix2Lock.RUnlock()
		return radix2Factors[input_len]
	}

	radix2Lock.RUnlock()
	radix2Lock.Lock()
	defer radix2Lock.Unlock()

	if !hasRadix2Factors(input_len) {
		for i, p := 8, 4; i <= input_len; i, p = i<<1, i {
			if radix2Factors[i] == nil {
				radix2Factors[i] = make([]complex128, i)

				for n, j := 0, 0; n < i; n, j = n+2, j+1 {
					radix2Factors[i][n] = radix2Factors[p][j]
				}

				for n := 1; n < i; n += 2 {
					sin, cos := math.Sincos(-2 * math.Pi / float64(i) * float64(n))
					radix2Factors[i][n] = complex(cos, sin)
				}
			}
		}
	}

	return radix2Factors[input_len]
}

func hasRadix2Factors(idx int) bool {
	return radix2Factors[idx] != nil
}

func getBluesteinFactors(input_len int) ([]complex128, []complex128) {
	bluesteinLock.RLock()

	if hasBluesteinFactors(input_len) {
		defer bluesteinLock.RUnlock()
		return bluesteinFactors[input_len], bluesteinInvFactors[input_len]
	}

	bluesteinLock.RUnlock()
	bluesteinLock.Lock()
	defer bluesteinLock.Unlock()

	if !hasBluesteinFactors(input_len) {
		bluesteinFactors[input_len] = make([]complex128, input_len)
		bluesteinInvFactors[input_len] = make([]complex128, input_len)

		var sin, cos float64
		for i := 0; i < input_len; i++ {
			if i == 0 {
				sin, cos = 0, 1
			} else {
				sin, cos = math.Sincos(math.Pi / float64(input_len) * float64(i*i))
			}
			bluesteinFactors[input_len][i] = complex(cos, sin)
			bluesteinInvFactors[input_len][i] = complex(cos, -sin)
		}
	}

	return bluesteinFactors[input_len], bluesteinInvFactors[input_len]
}

func hasBluesteinFactors(idx int) bool {
	return bluesteinFactors[idx] != nil
}

// FFTReal returns the forward FFT of the real-valued slice.
func FFTReal(x []float64) []complex128 {
	return FFT(dsputils.ToComplex(x))
}

// IFFTReal returns the inverse FFT of the real-valued slice.
func IFFTReal(x []float64) []complex128 {
	return IFFT(dsputils.ToComplex(x))
}

// IFFT returns the inverse FFT of the complex-valued slice.
func IFFT(x []complex128) []complex128 {
	lx := len(x)
	r := make([]complex128, lx)

	// Reverse inputs, which is calculated with modulo N, hence x[0] as an outlier
	r[0] = x[0]
	for i := 1; i < lx; i++ {
		r[i] = x[lx-i]
	}

	r = FFT(r)

	N := complex(float64(lx), 0)
	for n, _ := range r {
		r[n] /= N
	}
	return r
}

// Convolve returns the convolution of x * y.
func Convolve(x, y []complex128) []complex128 {
	if len(x) != len(y) {
		panic("arrays not of equal size")
	}

	fft_x := FFT(x)
	fft_y := FFT(y)

	r := make([]complex128, len(x))
	for i := 0; i < len(r); i++ {
		r[i] = fft_x[i] * fft_y[i]
	}

	return IFFT(r)
}

// FFT returns the forward FFT of the complex-valued slice.
func FFT(x []complex128) []complex128 {
	lx := len(x)

	// todo: non-hack handling length <= 1 cases
	if lx <= 1 {
		r := make([]complex128, lx)
		copy(r, x)
		return r
	}

	if dsputils.IsPowerOf2(lx) {
		return radix2FFT(x)
	}

	return bluesteinFFT(x)
}

var (
	worker_pool_size = 0
)

// SetWorkerPoolSize sets the number of workers during FFT computation on multicore systems.
// If n is 0 (the default), then GOMAXPROCS workers will be created.
func SetWorkerPoolSize(n int) {
	if n < 0 {
		n = 0
	}

	worker_pool_size = n
}

type fft_work struct {
	start, end int
}

// radix2FFT returns the FFT calculated using the radix-2 DIT Cooley-Tukey algorithm.
func radix2FFT(x []complex128) []complex128 {
	lx := len(x)
	factors := getRadix2Factors(lx)

	t := make([]complex128, lx) // temp
	r := reorderData(x)

	var blocks, stage, s_2 int

	jobs := make(chan *fft_work, lx)
	results := make(chan bool, lx)

	num_workers := worker_pool_size
	if (num_workers) == 0 {
		num_workers = runtime.GOMAXPROCS(0)
	}

	idx_diff := lx / num_workers
	if idx_diff < 2 {
		idx_diff = 2
	}

	worker := func() {
		for work := range jobs {
			for nb := work.start; nb < work.end; nb += stage {
				if stage != 2 {
					for j := 0; j < s_2; j++ {
						idx := j + nb
						idx2 := idx + s_2
						ridx := r[idx]
						w_n := r[idx2] * factors[blocks*j]
						t[idx] = ridx + w_n
						t[idx2] = ridx - w_n
					}
				} else {
					n1 := nb + 1
					rn := r[nb]
					rn1 := r[n1]
					t[nb] = rn + rn1
					t[n1] = rn - rn1
				}
			}

			results <- true
		}
	}

	for i := 0; i < num_workers; i++ {
		go worker()
	}
	defer close(jobs)

	for stage = 2; stage <= lx; stage <<= 1 {
		blocks = lx / stage
		s_2 = stage / 2
		workers_spawned := 0

		for start, end := 0, stage; ; {
			if end-start >= idx_diff || end == lx {
				workers_spawned++
				jobs <- &fft_work{start, end}

				if end == lx {
					break
				}

				start = end
			}

			end += stage
		}

		for n := 0; n < workers_spawned; n++ {
			<-results
		}

		r, t = t, r
	}

	return r
}

// reorderData returns a copy of x reordered for the DFT.
func reorderData(x []complex128) []complex128 {
	lx := uint(len(x))
	r := make([]complex128, lx)
	s := log2(lx)

	var n uint
	for ; n < lx; n++ {
		r[reverseBits(n, s)] = x[n]
	}

	return r
}

// log2 returns the log base 2 of v
// from: http://graphics.stanford.edu/~seander/bithacks.html#IntegerLogObvious
func log2(v uint) uint {
	var r uint

	for v >>= 1; v != 0; v >>= 1 {
		r++
	}

	return r
}

// reverseBits returns the first s bits of v in reverse order
// from: http://graphics.stanford.edu/~seander/bithacks.html#BitReverseObvious
func reverseBits(v, s uint) uint {
	var r uint

	// Since we aren't reversing all the bits in v (just the first s bits),
	// we only need the first bit of v instead of a full copy.
	r = v & 1
	s--

	for v >>= 1; v != 0; v >>= 1 {
		r <<= 1
		r |= v & 1
		s--
	}

	return r << s
}

// bluesteinFFT returns the FFT calculated using the Bluestein algorithm.
func bluesteinFFT(x []complex128) []complex128 {
	lx := len(x)
	a := dsputils.ZeroPad(x, dsputils.NextPowerOf2(lx*2-1))
	la := len(a)
	factors, invFactors := getBluesteinFactors(lx)

	for n, v := range x {
		a[n] = v * invFactors[n]
	}

	b := make([]complex128, la)
	for i := 0; i < lx; i++ {
		b[i] = factors[i]

		if i != 0 {
			b[la-i] = factors[i]
		}
	}

	r := Convolve(a, b)

	for i := 0; i < lx; i++ {
		r[i] *= invFactors[i]
	}

	return r[:lx]
}

// FFT2Real returns the 2-dimensional, forward FFT of the real-valued matrix.
func FFT2Real(x [][]float64) [][]complex128 {
	return FFT2(dsputils.ToComplex2(x))
}

// FFT2 returns the 2-dimensional, forward FFT of the complex-valued matrix.
func FFT2(x [][]complex128) [][]complex128 {
	return computeFFT2(x, FFT)
}

// IFFT2Real returns the 2-dimensional, inverse FFT of the real-valued matrix.
func IFFT2Real(x [][]float64) [][]complex128 {
	return IFFT2(dsputils.ToComplex2(x))
}

// IFFT2 returns the 2-dimensional, inverse FFT of the complex-valued matrix.
func IFFT2(x [][]complex128) [][]complex128 {
	return computeFFT2(x, IFFT)
}

func computeFFT2(x [][]complex128, fftFunc func([]complex128) []complex128) [][]complex128 {
	rows := len(x)
	if rows == 0 {
		panic("empty input array")
	}

	cols := len(x[0])
	r := make([][]complex128, rows)
	for i := 0; i < rows; i++ {
		if len(x[i]) != cols {
			panic("ragged input array")
		}
		r[i] = make([]complex128, cols)
	}

	for i := 0; i < cols; i++ {
		t := make([]complex128, rows)
		for j := 0; j < rows; j++ {
			t[j] = x[j][i]
		}

		for n, v := range fftFunc(t) {
			r[n][i] = v
		}
	}

	for n, v := range r {
		r[n] = fftFunc(v)
	}

	return r
}

// FFTN returns the forward FFT of the matrix m, computed in all N dimensions.
func FFTN(m *dsputils.Matrix) *dsputils.Matrix {
	return computeFFTN(m, FFT)
}

// IFFTN returns the forward FFT of the matrix m, computed in all N dimensions.
func IFFTN(m *dsputils.Matrix) *dsputils.Matrix {
	return computeFFTN(m, IFFT)
}

func computeFFTN(m *dsputils.Matrix, fftFunc func([]complex128) []complex128) *dsputils.Matrix {
	dims := m.Dimensions()
	t := m.Copy()
	r := dsputils.MakeEmptyMatrix(dims)

	for n := range dims {
		dims[n] -= 1
	}

	for n := range dims {
		d := make([]int, len(dims))
		copy(d, dims)
		d[n] = -1

		for {
			r.SetDim(fftFunc(t.Dim(d)), d)

			if !decrDim(d, dims) {
				break
			}
		}

		r, t = t, r
	}

	return t
}

// decrDim decrements an element of x by 1, skipping all -1s, and wrapping up to d.
// If a value is 0, it will be reset to its correspending value in d, and will carry one from the next non -1 value to the right.
// Returns true if decremented, else false.
func decrDim(x, d []int) bool {
	for n, v := range x {
		if v == -1 {
			continue
		} else if v == 0 {
			i := n
			// find the next element to decrement
			for ; i < len(x); i++ {
				if x[i] == -1 {
					continue
				} else if x[i] == 0 {
					x[i] = d[i]
				} else {
					x[i] -= 1
					return true
				}
			}

			// no decrement
			return false
		} else {
			x[n] -= 1
			return true
		}
	}

	return false
}
<commit_msg>Cooler * symbol for convolve<commit_after>/*
 * Copyright (c) 2011 Matt Jibson <matt.jibson@gmail.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

// Package fft provides forward and inverse fast Fourier transform functions.
package fft

import (
	"math"
	"runtime"
	"sync"

	"github.com/mjibson/go-dsp/dsputils"
)

var (
	radix2Lock    sync.RWMutex
	radix2Factors = map[int][]complex128{
		4: {complex(1, 0), complex(0, -1), complex(-1, 0), complex(0, 1)},
	}
)

var (
	bluesteinLock       sync.RWMutex
	bluesteinFactors    = map[int][]complex128{}
	bluesteinInvFactors = map[int][]complex128{}
)

// EnsureRadix2Factors ensures that all radix 2 factors are computed for inputs
// of length input_len. This is used to precompute needed factors for known
// sizes. Generally should only be used for benchmarks.
func EnsureRadix2Factors(input_len int) {
	getRadix2Factors(input_len)
}

func getRadix2Factors(input_len int) []complex128 {
	radix2Lock.RLock()

	if hasRadix2Factors(input_len) {
		defer radix2Lock.RUnlock()
		return radix2Factors[input_len]
	}

	radix2Lock.RUnlock()
	radix2Lock.Lock()
	defer radix2Lock.Unlock()

	if !hasRadix2Factors(input_len) {
		for i, p := 8, 4; i <= input_len; i, p = i<<1, i {
			if radix2Factors[i] == nil {
				radix2Factors[i] = make([]complex128, i)

				for n, j := 0, 0; n < i; n, j = n+2, j+1 {
					radix2Factors[i][n] = radix2Factors[p][j]
				}

				for n := 1; n < i; n += 2 {
					sin, cos := math.Sincos(-2 * math.Pi / float64(i) * float64(n))
					radix2Factors[i][n] = complex(cos, sin)
				}
			}
		}
	}

	return radix2Factors[input_len]
}

func hasRadix2Factors(idx int) bool {
	return radix2Factors[idx] != nil
}

func getBluesteinFactors(input_len int) ([]complex128, []complex128) {
	bluesteinLock.RLock()

	if hasBluesteinFactors(input_len) {
		defer bluesteinLock.RUnlock()
		return bluesteinFactors[input_len], bluesteinInvFactors[input_len]
	}

	bluesteinLock.RUnlock()
	bluesteinLock.Lock()
	defer bluesteinLock.Unlock()

	if !hasBluesteinFactors(input_len) {
		bluesteinFactors[input_len] = make([]complex128, input_len)
		bluesteinInvFactors[input_len] = make([]complex128, input_len)

		var sin, cos float64
		for i := 0; i < input_len; i++ {
			if i == 0 {
				sin, cos = 0, 1
			} else {
				sin, cos = math.Sincos(math.Pi / float64(input_len) * float64(i*i))
			}
			bluesteinFactors[input_len][i] = complex(cos, sin)
			bluesteinInvFactors[input_len][i] = complex(cos, -sin)
		}
	}

	return bluesteinFactors[input_len], bluesteinInvFactors[input_len]
}

func hasBluesteinFactors(idx int) bool {
	return bluesteinFactors[idx] != nil
}

// FFTReal returns the forward FFT of the real-valued slice.
func FFTReal(x []float64) []complex128 {
	return FFT(dsputils.ToComplex(x))
}

// IFFTReal returns the inverse FFT of the real-valued slice.
func IFFTReal(x []float64) []complex128 {
	return IFFT(dsputils.ToComplex(x))
}

// IFFT returns the inverse FFT of the complex-valued slice.
func IFFT(x []complex128) []complex128 {
	lx := len(x)
	r := make([]complex128, lx)

	// Reverse inputs, which is calculated with modulo N, hence x[0] as an outlier
	r[0] = x[0]
	for i := 1; i < lx; i++ {
		r[i] = x[lx-i]
	}

	r = FFT(r)

	N := complex(float64(lx), 0)
	for n, _ := range r {
		r[n] /= N
	}
	return r
}

// Convolve returns the convolution of x ∗ y.
func Convolve(x, y []complex128) []complex128 {
	if len(x) != len(y) {
		panic("arrays not of equal size")
	}

	fft_x := FFT(x)
	fft_y := FFT(y)

	r := make([]complex128, len(x))
	for i := 0; i < len(r); i++ {
		r[i] = fft_x[i] * fft_y[i]
	}

	return IFFT(r)
}

// FFT returns the forward FFT of the complex-valued slice.
func FFT(x []complex128) []complex128 {
	lx := len(x)

	// todo: non-hack handling length <= 1 cases
	if lx <= 1 {
		r := make([]complex128, lx)
		copy(r, x)
		return r
	}

	if dsputils.IsPowerOf2(lx) {
		return radix2FFT(x)
	}

	return bluesteinFFT(x)
}

var (
	worker_pool_size = 0
)

// SetWorkerPoolSize sets the number of workers during FFT computation on multicore systems.
// If n is 0 (the default), then GOMAXPROCS workers will be created.
func SetWorkerPoolSize(n int) {
	if n < 0 {
		n = 0
	}

	worker_pool_size = n
}

type fft_work struct {
	start, end int
}

// radix2FFT returns the FFT calculated using the radix-2 DIT Cooley-Tukey algorithm.
func radix2FFT(x []complex128) []complex128 {
	lx := len(x)
	factors := getRadix2Factors(lx)

	t := make([]complex128, lx) // temp
	r := reorderData(x)

	var blocks, stage, s_2 int

	jobs := make(chan *fft_work, lx)
	results := make(chan bool, lx)

	num_workers := worker_pool_size
	if (num_workers) == 0 {
		num_workers = runtime.GOMAXPROCS(0)
	}

	idx_diff := lx / num_workers
	if idx_diff < 2 {
		idx_diff = 2
	}

	worker := func() {
		for work := range jobs {
			for nb := work.start; nb < work.end; nb += stage {
				if stage != 2 {
					for j := 0; j < s_2; j++ {
						idx := j + nb
						idx2 := idx + s_2
						ridx := r[idx]
						w_n := r[idx2] * factors[blocks*j]
						t[idx] = ridx + w_n
						t[idx2] = ridx - w_n
					}
				} else {
					n1 := nb + 1
					rn := r[nb]
					rn1 := r[n1]
					t[nb] = rn + rn1
					t[n1] = rn - rn1
				}
			}

			results <- true
		}
	}

	for i := 0; i < num_workers; i++ {
		go worker()
	}
	defer close(jobs)

	for stage = 2; stage <= lx; stage <<= 1 {
		blocks = lx / stage
		s_2 = stage / 2
		workers_spawned := 0

		for start, end := 0, stage; ; {
			if end-start >= idx_diff || end == lx {
				workers_spawned++
				jobs <- &fft_work{start, end}

				if end == lx {
					break
				}

				start = end
			}

			end += stage
		}

		for n := 0; n < workers_spawned; n++ {
			<-results
		}

		r, t = t, r
	}

	return r
}

// reorderData returns a copy of x reordered for the DFT.
func reorderData(x []complex128) []complex128 {
	lx := uint(len(x))
	r := make([]complex128, lx)
	s := log2(lx)

	var n uint
	for ; n < lx; n++ {
		r[reverseBits(n, s)] = x[n]
	}

	return r
}

// log2 returns the log base 2 of v
// from: http://graphics.stanford.edu/~seander/bithacks.html#IntegerLogObvious
func log2(v uint) uint {
	var r uint

	for v >>= 1; v != 0; v >>= 1 {
		r++
	}

	return r
}

// reverseBits returns the first s bits of v in reverse order
// from: http://graphics.stanford.edu/~seander/bithacks.html#BitReverseObvious
func reverseBits(v, s uint) uint {
	var r uint

	// Since we aren't reversing all the bits in v (just the first s bits),
	// we only need the first bit of v instead of a full copy.
	r = v & 1
	s--

	for v >>= 1; v != 0; v >>= 1 {
		r <<= 1
		r |= v & 1
		s--
	}

	return r << s
}

// bluesteinFFT returns the FFT calculated using the Bluestein algorithm.
func bluesteinFFT(x []complex128) []complex128 {
	lx := len(x)
	a := dsputils.ZeroPad(x, dsputils.NextPowerOf2(lx*2-1))
	la := len(a)
	factors, invFactors := getBluesteinFactors(lx)

	for n, v := range x {
		a[n] = v * invFactors[n]
	}

	b := make([]complex128, la)
	for i := 0; i < lx; i++ {
		b[i] = factors[i]

		if i != 0 {
			b[la-i] = factors[i]
		}
	}

	r := Convolve(a, b)

	for i := 0; i < lx; i++ {
		r[i] *= invFactors[i]
	}

	return r[:lx]
}

// FFT2Real returns the 2-dimensional, forward FFT of the real-valued matrix.
func FFT2Real(x [][]float64) [][]complex128 {
	return FFT2(dsputils.ToComplex2(x))
}

// FFT2 returns the 2-dimensional, forward FFT of the complex-valued matrix.
func FFT2(x [][]complex128) [][]complex128 {
	return computeFFT2(x, FFT)
}

// IFFT2Real returns the 2-dimensional, inverse FFT of the real-valued matrix.
func IFFT2Real(x [][]float64) [][]complex128 {
	return IFFT2(dsputils.ToComplex2(x))
}

// IFFT2 returns the 2-dimensional, inverse FFT of the complex-valued matrix.
func IFFT2(x [][]complex128) [][]complex128 {
	return computeFFT2(x, IFFT)
}

func computeFFT2(x [][]complex128, fftFunc func([]complex128) []complex128) [][]complex128 {
	rows := len(x)
	if rows == 0 {
		panic("empty input array")
	}

	cols := len(x[0])
	r := make([][]complex128, rows)
	for i := 0; i < rows; i++ {
		if len(x[i]) != cols {
			panic("ragged input array")
		}
		r[i] = make([]complex128, cols)
	}

	for i := 0; i < cols; i++ {
		t := make([]complex128, rows)
		for j := 0; j < rows; j++ {
			t[j] = x[j][i]
		}

		for n, v := range fftFunc(t) {
			r[n][i] = v
		}
	}

	for n, v := range r {
		r[n] = fftFunc(v)
	}

	return r
}

// FFTN returns the forward FFT of the matrix m, computed in all N dimensions.
func FFTN(m *dsputils.Matrix) *dsputils.Matrix {
	return computeFFTN(m, FFT)
}

// IFFTN returns the forward FFT of the matrix m, computed in all N dimensions.
func IFFTN(m *dsputils.Matrix) *dsputils.Matrix {
	return computeFFTN(m, IFFT)
}

func computeFFTN(m *dsputils.Matrix, fftFunc func([]complex128) []complex128) *dsputils.Matrix {
	dims := m.Dimensions()
	t := m.Copy()
	r := dsputils.MakeEmptyMatrix(dims)

	for n := range dims {
		dims[n] -= 1
	}

	for n := range dims {
		d := make([]int, len(dims))
		copy(d, dims)
		d[n] = -1

		for {
			r.SetDim(fftFunc(t.Dim(d)), d)

			if !decrDim(d, dims) {
				break
			}
		}

		r, t = t, r
	}

	return t
}

// decrDim decrements an element of x by 1, skipping all -1s, and wrapping up to d.
// If a value is 0, it will be reset to its correspending value in d, and will carry one from the next non -1 value to the right.
// Returns true if decremented, else false.
func decrDim(x, d []int) bool {
	for n, v := range x {
		if v == -1 {
			continue
		} else if v == 0 {
			i := n
			// find the next element to decrement
			for ; i < len(x); i++ {
				if x[i] == -1 {
					continue
				} else if x[i] == 0 {
					x[i] = d[i]
				} else {
					x[i] -= 1
					return true
				}
			}

			// no decrement
			return false
		} else {
			x[n] -= 1
			return true
		}
	}

	return false
}
<|endoftext|>


<commit_before>// Copyright (C) 2010, Kyle Lemons <kyle@kylelemons.net>.  All rights reserved.

package log4go

import (
	"os"
	"fmt"
	"time"
)

type FileFactory func(string) (*os.File, error)

// This log writer sends output to a file
type FileLogWriter struct {
	rec chan *LogRecord
	rot chan bool

	// The opened file
	filename string
	file     *os.File
	filefactory FileFactory

	// The logging format
	format string

	// File header/trailer
	header, trailer string

	// Rotate at linecount
	maxlines          int
	maxlines_curlines int

	// Rotate at size
	maxsize         int
	maxsize_cursize int

	// Rotate daily
	daily          bool
	daily_opendate int

	// Keep old logfiles (.001, .002, etc)
	rotate bool
}

// This is the FileLogWriter's output method
func (w *FileLogWriter) LogWrite(rec *LogRecord) {
	w.rec <- rec
}

func (w *FileLogWriter) Close() {
	close(w.rec)
}

func internalFileFactory(fn string) (*os.File, error) {
	return os.OpenFile(fn, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0660)
}

// NewFileLogWriter creates a new LogWriter which writes to the given file and
// has rotation enabled if rotate is true.
//
// If rotate is true, any time a new log file is opened, the old one is renamed
// with a .### extension to preserve it.  The various Set* methods can be used
// to configure log rotation based on lines, size, and daily.
//
// The standard log-line format is:
//   [%D %T] [%L] (%S) %M
func NewFileLogWriter(fname string, rotate bool) *FileLogWriter {
	w := &FileLogWriter{
		rec:      		make(chan *LogRecord, LogBufferLength),
		rot:      		make(chan bool),
		filename: 		fname,
		format:   		"[%D %T] [%L] (%S) %M",
		rotate:   		rotate,
		filefactory: 	internalFileFactory,
	}

	// open the file for the first time
	if err := w.intRotate(); err != nil {
		fmt.Fprintf(os.Stderr, "FileLogWriter(%q): %s\n", w.filename, err)
		return nil
	}

	go func() {
		defer func() {
			if w.file != nil {
				fmt.Fprint(w.file, FormatLogRecord(w.trailer, &LogRecord{Created: time.Now()}))
				w.file.Close()
			}
		}()

		for {
			select {
			case <-w.rot:
				if err := w.intRotate(); err != nil {
					fmt.Fprintf(os.Stderr, "FileLogWriter(%q): %s\n", w.filename, err)
					return
				}
			case rec, ok := <-w.rec:
				if !ok {
					return
				}
				now := time.Now()
				if (w.maxlines > 0 && w.maxlines_curlines >= w.maxlines) ||
					(w.maxsize > 0 && w.maxsize_cursize >= w.maxsize) ||
					(w.daily && now.Day() != w.daily_opendate) {
					if err := w.intRotate(); err != nil {
						fmt.Fprintf(os.Stderr, "FileLogWriter(%q): %s\n", w.filename, err)
						return
					}
				}

				// Perform the write
				n, err := fmt.Fprint(w.file, FormatLogRecord(w.format, rec))
				if err != nil {
					fmt.Fprintf(os.Stderr, "FileLogWriter(%q): %s\n", w.filename, err)
					return
				}

				// Update the counts
				w.maxlines_curlines++
				w.maxsize_cursize += n
			}
		}
	}()

	return w
}

// Request that the logs rotate
func (w *FileLogWriter) Rotate() {
	w.rot <- true
}

// If this is called in a threaded context, it MUST be synchronized
func (w *FileLogWriter) intRotate() error {
	// Close any log file that may be open
	if w.file != nil {
		fmt.Fprint(w.file, FormatLogRecord(w.trailer, &LogRecord{Created: time.Now()}))
		w.file.Close()
	}

	// If we are keeping log files, move it to the next available number
	if w.rotate {
		_, err := os.Lstat(w.filename)
		if err == nil { // file exists
			// Find the next available number
			num := 1
			fname := ""
			for ; err == nil && num <= 999; num++ {
				fname = w.filename + fmt.Sprintf(".%03d", num)
				_, err = os.Lstat(fname)
			}
			// return error if the last file checked still existed
			if err == nil {
				return fmt.Errorf("Rotate: Cannot find free log number to rename %s\n", w.filename)
			}

			// Rename the file to its newfound home
			err = os.Rename(w.filename, fname)
			if err != nil {
				return fmt.Errorf("Rotate: %s\n", err)
			}
		}
	}

	// Open the log file
	fd, err := w.filefactory(w.filename)
	if err != nil {
		return err
	}
	w.file = fd

	now := time.Now()
	fmt.Fprint(w.file, FormatLogRecord(w.header, &LogRecord{Created: now}))

	// Set the daily open date to the current date
	w.daily_opendate = now.Day()

	// initialize rotation values
	w.maxlines_curlines = 0
	w.maxsize_cursize = 0

	return nil
}

// Set the logging format (chainable).  Must be called before the first log
// message is written.
func (w *FileLogWriter) SetFormat(format string) *FileLogWriter {
	w.format = format
	return w
}

func (w *FileLogWriter) SetFileFactory(factory FileFactory) *FileLogWriter {
	w.filefactory = factory
	return w
}

// Set the logfile header and footer (chainable).  Must be called before the first log
// message is written.  These are formatted similar to the FormatLogRecord (e.g.
// you can use %D and %T in your header/footer for date and time).
func (w *FileLogWriter) SetHeadFoot(head, foot string) *FileLogWriter {
	w.header, w.trailer = head, foot
	if w.maxlines_curlines == 0 {
		fmt.Fprint(w.file, FormatLogRecord(w.header, &LogRecord{Created: time.Now()}))
	}
	return w
}

// Set rotate at linecount (chainable). Must be called before the first log
// message is written.
func (w *FileLogWriter) SetRotateLines(maxlines int) *FileLogWriter {
	//fmt.Fprintf(os.Stderr, "FileLogWriter.SetRotateLines: %v\n", maxlines)
	w.maxlines = maxlines
	return w
}

// Set rotate at size (chainable). Must be called before the first log message
// is written.
func (w *FileLogWriter) SetRotateSize(maxsize int) *FileLogWriter {
	//fmt.Fprintf(os.Stderr, "FileLogWriter.SetRotateSize: %v\n", maxsize)
	w.maxsize = maxsize
	return w
}

// Set rotate daily (chainable). Must be called before the first log message is
// written.
func (w *FileLogWriter) SetRotateDaily(daily bool) *FileLogWriter {
	//fmt.Fprintf(os.Stderr, "FileLogWriter.SetRotateDaily: %v\n", daily)
	w.daily = daily
	return w
}

// SetRotate changes whether or not the old logs are kept. (chainable) Must be
// called before the first log message is written.  If rotate is false, the
// files are overwritten; otherwise, they are rotated to another file before the
// new log is opened.
func (w *FileLogWriter) SetRotate(rotate bool) *FileLogWriter {
	//fmt.Fprintf(os.Stderr, "FileLogWriter.SetRotate: %v\n", rotate)
	w.rotate = rotate
	return w
}

// NewXMLLogWriter is a utility method for creating a FileLogWriter set up to
// output XML record log messages instead of line-based ones.
func NewXMLLogWriter(fname string, rotate bool) *FileLogWriter {
	return NewFileLogWriter(fname, rotate).SetFormat(
		`	<record level="%L">
		<timestamp>%D %T</timestamp>
		<source>%S</source>
		<message>%M</message>
	</record>`).SetHeadFoot("<log created=\"%D %T\">", "</log>")
}
<commit_msg>0660 fix<commit_after>// Copyright (C) 2010, Kyle Lemons <kyle@kylelemons.net>.  All rights reserved.

package log4go

import (
	"os"
	"fmt"
	"time"
)

type FileFactory func(string) (*os.File, error)

// This log writer sends output to a file
type FileLogWriter struct {
	rec chan *LogRecord
	rot chan bool

	// The opened file
	filename string
	file     *os.File
	filefactory FileFactory

	// The logging format
	format string

	// File header/trailer
	header, trailer string

	// Rotate at linecount
	maxlines          int
	maxlines_curlines int

	// Rotate at size
	maxsize         int
	maxsize_cursize int

	// Rotate daily
	daily          bool
	daily_opendate int

	// Keep old logfiles (.001, .002, etc)
	rotate bool
}

// This is the FileLogWriter's output method
func (w *FileLogWriter) LogWrite(rec *LogRecord) {
	w.rec <- rec
}

func (w *FileLogWriter) Close() {
	close(w.rec)
}

func DefaultFileFactory(fn string) (*os.File, error) {
	return os.OpenFile(fn, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0660)
}

// NewFileLogWriter creates a new LogWriter which writes to the given file and
// has rotation enabled if rotate is true.
//
// If rotate is true, any time a new log file is opened, the old one is renamed
// with a .### extension to preserve it.  The various Set* methods can be used
// to configure log rotation based on lines, size, and daily.
//
// The standard log-line format is:
//   [%D %T] [%L] (%S) %M
func NewFileLogWriter(fname string, rotate bool, filefactory *FileFactory) *FileLogWriter {
	w := &FileLogWriter{
		rec:      		make(chan *LogRecord, LogBufferLength),
		rot:      		make(chan bool),
		filename: 		fname,
		format:   		"[%D %T] [%L] (%S) %M",
		rotate:   		rotate,
		filefactory: 	filefactory,
	}

	// open the file for the first time
	if err := w.intRotate(); err != nil {
		fmt.Fprintf(os.Stderr, "FileLogWriter(%q): %s\n", w.filename, err)
		return nil
	}

	go func() {
		defer func() {
			if w.file != nil {
				fmt.Fprint(w.file, FormatLogRecord(w.trailer, &LogRecord{Created: time.Now()}))
				w.file.Close()
			}
		}()

		for {
			select {
			case <-w.rot:
				if err := w.intRotate(); err != nil {
					fmt.Fprintf(os.Stderr, "FileLogWriter(%q): %s\n", w.filename, err)
					return
				}
			case rec, ok := <-w.rec:
				if !ok {
					return
				}
				now := time.Now()
				if (w.maxlines > 0 && w.maxlines_curlines >= w.maxlines) ||
					(w.maxsize > 0 && w.maxsize_cursize >= w.maxsize) ||
					(w.daily && now.Day() != w.daily_opendate) {
					if err := w.intRotate(); err != nil {
						fmt.Fprintf(os.Stderr, "FileLogWriter(%q): %s\n", w.filename, err)
						return
					}
				}

				// Perform the write
				n, err := fmt.Fprint(w.file, FormatLogRecord(w.format, rec))
				if err != nil {
					fmt.Fprintf(os.Stderr, "FileLogWriter(%q): %s\n", w.filename, err)
					return
				}

				// Update the counts
				w.maxlines_curlines++
				w.maxsize_cursize += n
			}
		}
	}()

	return w
}

// Request that the logs rotate
func (w *FileLogWriter) Rotate() {
	w.rot <- true
}

// If this is called in a threaded context, it MUST be synchronized
func (w *FileLogWriter) intRotate() error {
	// Close any log file that may be open
	if w.file != nil {
		fmt.Fprint(w.file, FormatLogRecord(w.trailer, &LogRecord{Created: time.Now()}))
		w.file.Close()
	}

	// If we are keeping log files, move it to the next available number
	if w.rotate {
		_, err := os.Lstat(w.filename)
		if err == nil { // file exists
			// Find the next available number
			num := 1
			fname := ""
			for ; err == nil && num <= 999; num++ {
				fname = w.filename + fmt.Sprintf(".%03d", num)
				_, err = os.Lstat(fname)
			}
			// return error if the last file checked still existed
			if err == nil {
				return fmt.Errorf("Rotate: Cannot find free log number to rename %s\n", w.filename)
			}

			// Rename the file to its newfound home
			err = os.Rename(w.filename, fname)
			if err != nil {
				return fmt.Errorf("Rotate: %s\n", err)
			}
		}
	}

	// Open the log file
	fd, err := w.filefactory(w.filename)
	if err != nil {
		return err
	}
	w.file = fd

	now := time.Now()
	fmt.Fprint(w.file, FormatLogRecord(w.header, &LogRecord{Created: now}))

	// Set the daily open date to the current date
	w.daily_opendate = now.Day()

	// initialize rotation values
	w.maxlines_curlines = 0
	w.maxsize_cursize = 0

	return nil
}

// Set the logging format (chainable).  Must be called before the first log
// message is written.
func (w *FileLogWriter) SetFormat(format string) *FileLogWriter {
	w.format = format
	return w
}

func (w *FileLogWriter) SetFileFactory(factory FileFactory) *FileLogWriter {
	w.filefactory = factory
	return w
}

// Set the logfile header and footer (chainable).  Must be called before the first log
// message is written.  These are formatted similar to the FormatLogRecord (e.g.
// you can use %D and %T in your header/footer for date and time).
func (w *FileLogWriter) SetHeadFoot(head, foot string) *FileLogWriter {
	w.header, w.trailer = head, foot
	if w.maxlines_curlines == 0 {
		fmt.Fprint(w.file, FormatLogRecord(w.header, &LogRecord{Created: time.Now()}))
	}
	return w
}

// Set rotate at linecount (chainable). Must be called before the first log
// message is written.
func (w *FileLogWriter) SetRotateLines(maxlines int) *FileLogWriter {
	//fmt.Fprintf(os.Stderr, "FileLogWriter.SetRotateLines: %v\n", maxlines)
	w.maxlines = maxlines
	return w
}

// Set rotate at size (chainable). Must be called before the first log message
// is written.
func (w *FileLogWriter) SetRotateSize(maxsize int) *FileLogWriter {
	//fmt.Fprintf(os.Stderr, "FileLogWriter.SetRotateSize: %v\n", maxsize)
	w.maxsize = maxsize
	return w
}

// Set rotate daily (chainable). Must be called before the first log message is
// written.
func (w *FileLogWriter) SetRotateDaily(daily bool) *FileLogWriter {
	//fmt.Fprintf(os.Stderr, "FileLogWriter.SetRotateDaily: %v\n", daily)
	w.daily = daily
	return w
}

// SetRotate changes whether or not the old logs are kept. (chainable) Must be
// called before the first log message is written.  If rotate is false, the
// files are overwritten; otherwise, they are rotated to another file before the
// new log is opened.
func (w *FileLogWriter) SetRotate(rotate bool) *FileLogWriter {
	//fmt.Fprintf(os.Stderr, "FileLogWriter.SetRotate: %v\n", rotate)
	w.rotate = rotate
	return w
}

// NewXMLLogWriter is a utility method for creating a FileLogWriter set up to
// output XML record log messages instead of line-based ones.
func NewXMLLogWriter(fname string, rotate bool) *FileLogWriter {
	return NewFileLogWriter(fname, rotate).SetFormat(
		`	<record level="%L">
		<timestamp>%D %T</timestamp>
		<source>%S</source>
		<message>%M</message>
	</record>`).SetHeadFoot("<log created=\"%D %T\">", "</log>")
}
<|endoftext|>


<commit_before>package cmd

import (
	"context"
	"fmt"
	"github.com/zrepl/zrepl/zfs"
	"sort"
	"time"
)

type IntervalAutosnap struct {
	task             *Task
	DatasetFilter    zfs.DatasetFilter
	Prefix           string
	SnapshotInterval time.Duration
}

func (a *IntervalAutosnap) filterFilesystems() (fss []*zfs.DatasetPath, stop bool) {
	a.task.Enter("filter_filesystems")
	defer a.task.Finish()
	fss, err := zfs.ZFSListMapping(a.DatasetFilter)
	stop = err != nil || len(fss) == 0
	if err != nil {
		a.task.Log().WithError(err).Error("cannot list datasets")
	}
	if len(fss) == 0 {
		a.task.Log().Warn("no filesystem matching filesystem filter")
	}
	return fss, stop
}

func (a *IntervalAutosnap) findSyncPoint(fss []*zfs.DatasetPath) (syncPoint time.Time, err error) {
	a.task.Enter("find_sync_point")
	defer a.task.Finish()
	type snapTime struct {
		ds   *zfs.DatasetPath
		time time.Time
	}

	snaptimes := make([]snapTime, 0, len(fss))

	now := time.Now()

	a.task.Log().Debug("examine filesystem state")
	for _, d := range fss {

		l := a.task.Log().WithField(logFSField, d.ToString())

		fsvs, err := zfs.ZFSListFilesystemVersions(d, NewTypedPrefixFilter(a.Prefix, zfs.Snapshot))
		if err != nil {
			l.WithError(err).Error("cannot list filesystem versions")
			continue
		}
		if len(fsvs) <= 0 {
			l.WithField("prefix", a.Prefix).Info("no filesystem versions with prefix")
			continue
		}

		// Sort versions by creation
		sort.SliceStable(fsvs, func(i, j int) bool {
			return fsvs[i].CreateTXG < fsvs[j].CreateTXG
		})

		latest := fsvs[len(fsvs)-1]
		l.WithField("creation", latest.Creation).
			Debug("found latest snapshot")

		since := now.Sub(latest.Creation)
		if since < 0 {
			l.WithField("snapshot", latest.Name).
				WithField("creation", latest.Creation).
				Error("snapshot is from the future")
			continue
		}
		next := now
		if since < a.SnapshotInterval {
			next = latest.Creation.Add(a.SnapshotInterval)
		}
		snaptimes = append(snaptimes, snapTime{d, next})
	}

	if len(snaptimes) == 0 {
		snaptimes = append(snaptimes, snapTime{nil, now})
	}

	sort.Slice(snaptimes, func(i, j int) bool {
		return snaptimes[i].time.Before(snaptimes[j].time)
	})

	return snaptimes[0].time, nil

}

func (a *IntervalAutosnap) waitForSyncPoint(ctx context.Context, syncPoint time.Time) {
	a.task.Enter("wait_sync_point")
	defer a.task.Finish()

	const LOG_TIME_FMT string = time.ANSIC

	a.task.Log().WithField("sync_point", syncPoint.Format(LOG_TIME_FMT)).
		Info("wait for sync point")

	select {
	case <-ctx.Done():
		a.task.Log().WithError(ctx.Err()).Info("context done")
		return
	case <-time.After(syncPoint.Sub(time.Now())):
	}
}

func (a *IntervalAutosnap) syncUpRun(ctx context.Context, didSnaps chan struct{}) (stop bool) {
	a.task.Enter("sync_up")
	defer a.task.Finish()

	fss, stop := a.filterFilesystems()
	if stop {
		return true
	}

	syncPoint, err := a.findSyncPoint(fss)
	if err != nil {
		return true
	}

	a.waitForSyncPoint(ctx, syncPoint)

	a.task.Log().Debug("snapshot all filesystems to enable further snaps in lockstep")
	a.doSnapshots(didSnaps)
	return false
}

func (a *IntervalAutosnap) Run(ctx context.Context, didSnaps chan struct{}) {

	if a.syncUpRun(ctx, didSnaps) {
		return
	}

	// task drops back to idle here

	a.task.Log().Debug("setting up ticker in SnapshotInterval")
	ticker := time.NewTicker(a.SnapshotInterval)
	for {
		select {
		case <-ctx.Done():
			ticker.Stop()
			a.task.Log().WithError(ctx.Err()).Info("context done")
			return

		case <-ticker.C:
			a.doSnapshots(didSnaps)
		}
	}

}

func (a *IntervalAutosnap) doSnapshots(didSnaps chan struct{}) {

	a.task.Enter("do_snapshots")
	defer a.task.Finish()

	// don't cache the result from previous run in case the user added
	// a new dataset in the meantime
	ds, stop := a.filterFilesystems()
	if stop {
		return
	}

	// TODO channel programs -> allow a little jitter?
	for _, d := range ds {
		suffix := time.Now().In(time.UTC).Format("20060102_150405_000")
		snapname := fmt.Sprintf("%s%s", a.Prefix, suffix)

		l := a.task.Log().WithField(logFSField, d.ToString()).
			WithField("snapname", snapname)

		l.Info("create snapshot")
		err := zfs.ZFSSnapshot(d, snapname, false)
		if err != nil {
			a.task.Log().WithError(err).Error("cannot create snapshot")
		}

		l.Info("create corresponding bookmark")
		err = zfs.ZFSBookmark(d, snapname, snapname)
		if err != nil {
			a.task.Log().WithError(err).Error("cannot create bookmark")
		}

	}

	select {
	case didSnaps <- struct{}{}:
	default:
		a.task.Log().Error("warning: callback channel is full, discarding")
	}

}
<commit_msg>autosnap: do not treat zero fs filter results as fatal<commit_after>package cmd

import (
	"context"
	"fmt"
	"github.com/zrepl/zrepl/zfs"
	"sort"
	"time"
)

type IntervalAutosnap struct {
	task             *Task
	DatasetFilter    zfs.DatasetFilter
	Prefix           string
	SnapshotInterval time.Duration
}

func (a *IntervalAutosnap) filterFilesystems() (fss []*zfs.DatasetPath, stop bool) {
	a.task.Enter("filter_filesystems")
	defer a.task.Finish()
	fss, err := zfs.ZFSListMapping(a.DatasetFilter)
	stop = err != nil
	if err != nil {
		a.task.Log().WithError(err).Error("cannot list datasets")
	}
	if len(fss) == 0 {
		a.task.Log().Warn("no filesystem matching filesystem filter")
	}
	return fss, stop
}

func (a *IntervalAutosnap) findSyncPoint(fss []*zfs.DatasetPath) (syncPoint time.Time, err error) {
	a.task.Enter("find_sync_point")
	defer a.task.Finish()
	type snapTime struct {
		ds   *zfs.DatasetPath
		time time.Time
	}

	if len(fss) == 0 {
		return time.Now(), nil
	}

	snaptimes := make([]snapTime, 0, len(fss))

	now := time.Now()

	a.task.Log().Debug("examine filesystem state")
	for _, d := range fss {

		l := a.task.Log().WithField(logFSField, d.ToString())

		fsvs, err := zfs.ZFSListFilesystemVersions(d, NewTypedPrefixFilter(a.Prefix, zfs.Snapshot))
		if err != nil {
			l.WithError(err).Error("cannot list filesystem versions")
			continue
		}
		if len(fsvs) <= 0 {
			l.WithField("prefix", a.Prefix).Info("no filesystem versions with prefix")
			continue
		}

		// Sort versions by creation
		sort.SliceStable(fsvs, func(i, j int) bool {
			return fsvs[i].CreateTXG < fsvs[j].CreateTXG
		})

		latest := fsvs[len(fsvs)-1]
		l.WithField("creation", latest.Creation).
			Debug("found latest snapshot")

		since := now.Sub(latest.Creation)
		if since < 0 {
			l.WithField("snapshot", latest.Name).
				WithField("creation", latest.Creation).
				Error("snapshot is from the future")
			continue
		}
		next := now
		if since < a.SnapshotInterval {
			next = latest.Creation.Add(a.SnapshotInterval)
		}
		snaptimes = append(snaptimes, snapTime{d, next})
	}

	if len(snaptimes) == 0 {
		snaptimes = append(snaptimes, snapTime{nil, now})
	}

	sort.Slice(snaptimes, func(i, j int) bool {
		return snaptimes[i].time.Before(snaptimes[j].time)
	})

	return snaptimes[0].time, nil

}

func (a *IntervalAutosnap) waitForSyncPoint(ctx context.Context, syncPoint time.Time) {
	a.task.Enter("wait_sync_point")
	defer a.task.Finish()

	const LOG_TIME_FMT string = time.ANSIC

	a.task.Log().WithField("sync_point", syncPoint.Format(LOG_TIME_FMT)).
		Info("wait for sync point")

	select {
	case <-ctx.Done():
		a.task.Log().WithError(ctx.Err()).Info("context done")
		return
	case <-time.After(syncPoint.Sub(time.Now())):
	}
}

func (a *IntervalAutosnap) syncUpRun(ctx context.Context, didSnaps chan struct{}) (stop bool) {
	a.task.Enter("sync_up")
	defer a.task.Finish()

	fss, stop := a.filterFilesystems()
	if stop {
		return true
	}

	syncPoint, err := a.findSyncPoint(fss)
	if err != nil {
		return true
	}

	a.waitForSyncPoint(ctx, syncPoint)

	a.task.Log().Debug("snapshot all filesystems to enable further snaps in lockstep")
	a.doSnapshots(didSnaps)
	return false
}

func (a *IntervalAutosnap) Run(ctx context.Context, didSnaps chan struct{}) {

	if a.syncUpRun(ctx, didSnaps) {
		a.task.Log().Error("stoppping autosnap after error in sync up")
		return
	}

	// task drops back to idle here

	a.task.Log().Debug("setting up ticker in SnapshotInterval")
	ticker := time.NewTicker(a.SnapshotInterval)
	for {
		select {
		case <-ctx.Done():
			ticker.Stop()
			a.task.Log().WithError(ctx.Err()).Info("context done")
			return

		case <-ticker.C:
			a.doSnapshots(didSnaps)
		}
	}

}

func (a *IntervalAutosnap) doSnapshots(didSnaps chan struct{}) {

	a.task.Enter("do_snapshots")
	defer a.task.Finish()

	// don't cache the result from previous run in case the user added
	// a new dataset in the meantime
	ds, stop := a.filterFilesystems()
	if stop {
		return
	}

	// TODO channel programs -> allow a little jitter?
	for _, d := range ds {
		suffix := time.Now().In(time.UTC).Format("20060102_150405_000")
		snapname := fmt.Sprintf("%s%s", a.Prefix, suffix)

		l := a.task.Log().WithField(logFSField, d.ToString()).
			WithField("snapname", snapname)

		l.Info("create snapshot")
		err := zfs.ZFSSnapshot(d, snapname, false)
		if err != nil {
			a.task.Log().WithError(err).Error("cannot create snapshot")
		}

		l.Info("create corresponding bookmark")
		err = zfs.ZFSBookmark(d, snapname, snapname)
		if err != nil {
			a.task.Log().WithError(err).Error("cannot create bookmark")
		}

	}

	select {
	case didSnaps <- struct{}{}:
	default:
		a.task.Log().Error("warning: callback channel is full, discarding")
	}

}
<|endoftext|>


<commit_before>// Copyright 2019 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// The fetch command runs a server that fetches modules from a proxy and writes
// them to the discovery database.
package main

import (
	"bufio"
	"context"
	"flag"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	cloudtasks "cloud.google.com/go/cloudtasks/apiv2"
	"cloud.google.com/go/errorreporting"
	"cloud.google.com/go/profiler"
	"github.com/go-redis/redis/v7"
	"golang.org/x/discovery/internal/config"
	"golang.org/x/discovery/internal/database"
	"golang.org/x/discovery/internal/dcensus"
	"golang.org/x/discovery/internal/etl"
	"golang.org/x/discovery/internal/index"
	"golang.org/x/discovery/internal/queue"

	"golang.org/x/discovery/internal/log"
	"golang.org/x/discovery/internal/middleware"
	"golang.org/x/discovery/internal/postgres"
	"golang.org/x/discovery/internal/proxy"

	"contrib.go.opencensus.io/integrations/ocsql"
)

var (
	timeout    = config.GetEnv("GO_DISCOVERY_ETL_TIMEOUT_MINUTES", "10")
	queueName  = config.GetEnv("GO_DISCOVERY_ETL_TASK_QUEUE", "dev-fetch-tasks")
	workers    = flag.Int("workers", 10, "number of concurrent requests to the fetch service, when running locally")
	staticPath = flag.String("static", "content/static", "path to folder containing static files served")
)

func main() {
	flag.Parse()

	ctx := context.Background()

	cfg, err := config.Init(ctx)
	if err != nil {
		log.Fatal(ctx, err)
	}
	cfg.Dump(os.Stderr)

	if cfg.UseProfiler {
		if err := profiler.Start(profiler.Config{}); err != nil {
			log.Fatalf(ctx, "profiler.Start: %v", err)
		}
	}

	readProxyRemoved(ctx)

	// Wrap the postgres driver with OpenCensus instrumentation.
	driverName, err := ocsql.Register("postgres", ocsql.WithAllTraceOptions())
	if err != nil {
		log.Fatalf(ctx, "unable to register the ocsql driver: %v\n", err)
	}
	ddb, err := database.Open(driverName, cfg.DBConnInfo())
	if err != nil {
		log.Fatalf(ctx, "database.Open: %v", err)
	}
	db := postgres.New(ddb)
	defer db.Close()

	populateExcluded(ctx, db)

	indexClient, err := index.New(cfg.IndexURL)
	if err != nil {
		log.Fatal(ctx, err)
	}
	proxyClient, err := proxy.New(cfg.ProxyURL)
	if err != nil {
		log.Fatal(ctx, err)
	}
	fetchQueue := newQueue(ctx, cfg, proxyClient, db)
	reportingClient := reportingClient(ctx, cfg)
	redisClient := getRedis(ctx, cfg)
	server, err := etl.NewServer(cfg, db, indexClient, proxyClient, redisClient, fetchQueue, reportingClient, *staticPath)
	if err != nil {
		log.Fatal(ctx, err)
	}
	router := dcensus.NewRouter(nil)
	server.Install(router.Handle)

	views := append(dcensus.ClientViews, dcensus.ServerViews...)
	if err := dcensus.Init(cfg, views...); err != nil {
		log.Fatal(ctx, err)
	}
	// We are not currently forwarding any ports on AppEngine, so serving debug
	// information is broken.
	if !cfg.OnAppEngine() {
		dcensusServer, err := dcensus.NewServer()
		if err != nil {
			log.Fatal(ctx, err)
		}
		go http.ListenAndServe(cfg.DebugAddr("localhost:8001"), dcensusServer)
	}

	handlerTimeout, err := strconv.Atoi(timeout)
	if err != nil {
		log.Fatalf(ctx, "strconv.Atoi(%q): %v", timeout, err)
	}
	requestLogger := logger(ctx, cfg)
	mw := middleware.Chain(
		middleware.RequestLog(requestLogger),
		middleware.Timeout(time.Duration(handlerTimeout)*time.Minute),
	)
	http.Handle("/", mw(router))

	addr := cfg.HostAddr("localhost:8000")
	log.Infof(ctx, "Listening on addr %s", addr)
	log.Fatal(ctx, http.ListenAndServe(addr, nil))
}

func newQueue(ctx context.Context, cfg *config.Config, proxyClient *proxy.Client, db *postgres.DB) queue.Queue {
	if !cfg.OnAppEngine() {
		return queue.NewInMemory(ctx, proxyClient, db, *workers, etl.FetchAndUpdateState)
	}
	client, err := cloudtasks.NewClient(ctx)
	if err != nil {
		log.Fatal(ctx, err)
	}
	return queue.NewGCP(cfg, client, queueName)
}

func getRedis(ctx context.Context, cfg *config.Config) *redis.Client {
	var dialTimeout time.Duration
	if dl, ok := ctx.Deadline(); ok {
		dialTimeout = time.Until(dl)
	}
	if cfg.RedisHAHost != "" {
		return redis.NewClient(&redis.Options{
			Addr:        cfg.RedisHAHost + ":" + cfg.RedisHAPort,
			DialTimeout: dialTimeout,
			// We update completions with one big pipeline, so we need long write
			// timeouts. ReadTimeout is increased only to be consistent with
			// WriteTimeout.
			WriteTimeout: 5 * time.Minute,
			ReadTimeout:  5 * time.Minute,
		})
	}
	return nil
}

func reportingClient(ctx context.Context, cfg *config.Config) *errorreporting.Client {
	if !cfg.OnAppEngine() {
		return nil
	}
	reporter, err := errorreporting.NewClient(ctx, cfg.ProjectID, errorreporting.Config{
		ServiceName: cfg.ServiceID,
		OnError: func(err error) {
			log.Errorf(ctx, "Error reporting failed: %v", err)
		},
	})
	if err != nil {
		log.Fatal(ctx, err)
	}
	return reporter
}

func logger(ctx context.Context, cfg *config.Config) middleware.Logger {
	if cfg.OnAppEngine() {
		logger, err := log.UseStackdriver(ctx, cfg, "etl-log")
		if err != nil {
			log.Fatal(ctx, err)
		}
		return logger
	}
	return middleware.LocalLogger{}
}

// Read a file of module versions that we should ignore because
// the are in the index but not stored in the proxy.
// Format of the file: each line is
//     module@version
func readProxyRemoved(ctx context.Context) {
	filename := config.GetEnv("GO_DISCOVERY_PROXY_REMOVED", "")
	if filename == "" {
		return
	}
	f, err := os.Open(filename)
	if err != nil {
		log.Fatal(ctx, err)
	}
	defer f.Close()
	scan := bufio.NewScanner(f)
	for scan.Scan() {
		etl.ProxyRemoved[strings.TrimSpace(scan.Text())] = true
	}
	if err := scan.Err(); err != nil {
		log.Fatalf(ctx, "scanning %s: %v", filename, err)
	}
	log.Infof(ctx, "read %d excluded module versions from %s", len(etl.ProxyRemoved), filename)
}

// excludedPrefixes is a list of excluded prefixes and the reasons for exclusion.
// This is a permanent record of exclusions, in case the DB gets wiped or corrupted.
var excludedPrefixes = []struct {
	prefix, reason string
}{
	{
		"github.com/xvrzhao/site-monitor",
		"author requested https://groups.google.com/a/google.com/d/msg/go-discovery-feedback/oYtPw2Ob0fY/xxGikZK1AQAJ",
	},
	{
		"gioui.org/ui",
		"author requested https://groups.google.com/a/google.com/d/msg/go-discovery-feedback/CeMEn2E1zwo/q5S8HPn6BgAJ",
	},
	{
		"github.com/kortschak/unlicensable",
		"https://groups.google.com/g/golang-dev/c/mfiPCtJ1BGU/m/HDb3--vMEwAJk",
	},
}

// populateExcluded adds each element of excludedPrefixes to the excluded_prefixes
// table if it isn't already present.
func populateExcluded(ctx context.Context, db *postgres.DB) {
	user := os.Getenv("USER")
	if user == "" {
		user = "etl"
	}
	for _, ep := range excludedPrefixes {
		present, err := db.IsExcluded(ctx, ep.prefix)
		if err != nil {
			log.Fatalf(ctx, "db.IsExcluded(%q): %v", ep.prefix, err)
		}
		if !present {
			if err := db.InsertExcludedPrefix(ctx, ep.prefix, user, ep.reason); err != nil {
				log.Fatalf(ctx, "db.InsertExcludedPrefix(%q, %q, %q): %v", ep.prefix, user, ep.reason, err)
			}
		}
	}
}
<commit_msg>cmd/etl: move excluded prefixes to a file<commit_after>// Copyright 2019 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// The fetch command runs a server that fetches modules from a proxy and writes
// them to the discovery database.
package main

import (
	"bufio"
	"context"
	"flag"
	"fmt"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	cloudtasks "cloud.google.com/go/cloudtasks/apiv2"
	"cloud.google.com/go/errorreporting"
	"cloud.google.com/go/profiler"
	"github.com/go-redis/redis/v7"
	"golang.org/x/discovery/internal/config"
	"golang.org/x/discovery/internal/database"
	"golang.org/x/discovery/internal/dcensus"
	"golang.org/x/discovery/internal/etl"
	"golang.org/x/discovery/internal/index"
	"golang.org/x/discovery/internal/queue"

	"golang.org/x/discovery/internal/log"
	"golang.org/x/discovery/internal/middleware"
	"golang.org/x/discovery/internal/postgres"
	"golang.org/x/discovery/internal/proxy"

	"contrib.go.opencensus.io/integrations/ocsql"
)

var (
	timeout    = config.GetEnv("GO_DISCOVERY_ETL_TIMEOUT_MINUTES", "10")
	queueName  = config.GetEnv("GO_DISCOVERY_ETL_TASK_QUEUE", "dev-fetch-tasks")
	workers    = flag.Int("workers", 10, "number of concurrent requests to the fetch service, when running locally")
	staticPath = flag.String("static", "content/static", "path to folder containing static files served")
)

func main() {
	flag.Parse()

	ctx := context.Background()

	cfg, err := config.Init(ctx)
	if err != nil {
		log.Fatal(ctx, err)
	}
	cfg.Dump(os.Stderr)

	if cfg.UseProfiler {
		if err := profiler.Start(profiler.Config{}); err != nil {
			log.Fatalf(ctx, "profiler.Start: %v", err)
		}
	}

	readProxyRemoved(ctx)

	// Wrap the postgres driver with OpenCensus instrumentation.
	driverName, err := ocsql.Register("postgres", ocsql.WithAllTraceOptions())
	if err != nil {
		log.Fatalf(ctx, "unable to register the ocsql driver: %v\n", err)
	}
	ddb, err := database.Open(driverName, cfg.DBConnInfo())
	if err != nil {
		log.Fatalf(ctx, "database.Open: %v", err)
	}
	db := postgres.New(ddb)
	defer db.Close()

	populateExcluded(ctx, db)

	indexClient, err := index.New(cfg.IndexURL)
	if err != nil {
		log.Fatal(ctx, err)
	}
	proxyClient, err := proxy.New(cfg.ProxyURL)
	if err != nil {
		log.Fatal(ctx, err)
	}
	fetchQueue := newQueue(ctx, cfg, proxyClient, db)
	reportingClient := reportingClient(ctx, cfg)
	redisClient := getRedis(ctx, cfg)
	server, err := etl.NewServer(cfg, db, indexClient, proxyClient, redisClient, fetchQueue, reportingClient, *staticPath)
	if err != nil {
		log.Fatal(ctx, err)
	}
	router := dcensus.NewRouter(nil)
	server.Install(router.Handle)

	views := append(dcensus.ClientViews, dcensus.ServerViews...)
	if err := dcensus.Init(cfg, views...); err != nil {
		log.Fatal(ctx, err)
	}
	// We are not currently forwarding any ports on AppEngine, so serving debug
	// information is broken.
	if !cfg.OnAppEngine() {
		dcensusServer, err := dcensus.NewServer()
		if err != nil {
			log.Fatal(ctx, err)
		}
		go http.ListenAndServe(cfg.DebugAddr("localhost:8001"), dcensusServer)
	}

	handlerTimeout, err := strconv.Atoi(timeout)
	if err != nil {
		log.Fatalf(ctx, "strconv.Atoi(%q): %v", timeout, err)
	}
	requestLogger := logger(ctx, cfg)
	mw := middleware.Chain(
		middleware.RequestLog(requestLogger),
		middleware.Timeout(time.Duration(handlerTimeout)*time.Minute),
	)
	http.Handle("/", mw(router))

	addr := cfg.HostAddr("localhost:8000")
	log.Infof(ctx, "Listening on addr %s", addr)
	log.Fatal(ctx, http.ListenAndServe(addr, nil))
}

func newQueue(ctx context.Context, cfg *config.Config, proxyClient *proxy.Client, db *postgres.DB) queue.Queue {
	if !cfg.OnAppEngine() {
		return queue.NewInMemory(ctx, proxyClient, db, *workers, etl.FetchAndUpdateState)
	}
	client, err := cloudtasks.NewClient(ctx)
	if err != nil {
		log.Fatal(ctx, err)
	}
	return queue.NewGCP(cfg, client, queueName)
}

func getRedis(ctx context.Context, cfg *config.Config) *redis.Client {
	var dialTimeout time.Duration
	if dl, ok := ctx.Deadline(); ok {
		dialTimeout = time.Until(dl)
	}
	if cfg.RedisHAHost != "" {
		return redis.NewClient(&redis.Options{
			Addr:        cfg.RedisHAHost + ":" + cfg.RedisHAPort,
			DialTimeout: dialTimeout,
			// We update completions with one big pipeline, so we need long write
			// timeouts. ReadTimeout is increased only to be consistent with
			// WriteTimeout.
			WriteTimeout: 5 * time.Minute,
			ReadTimeout:  5 * time.Minute,
		})
	}
	return nil
}

func reportingClient(ctx context.Context, cfg *config.Config) *errorreporting.Client {
	if !cfg.OnAppEngine() {
		return nil
	}
	reporter, err := errorreporting.NewClient(ctx, cfg.ProjectID, errorreporting.Config{
		ServiceName: cfg.ServiceID,
		OnError: func(err error) {
			log.Errorf(ctx, "Error reporting failed: %v", err)
		},
	})
	if err != nil {
		log.Fatal(ctx, err)
	}
	return reporter
}

func logger(ctx context.Context, cfg *config.Config) middleware.Logger {
	if cfg.OnAppEngine() {
		logger, err := log.UseStackdriver(ctx, cfg, "etl-log")
		if err != nil {
			log.Fatal(ctx, err)
		}
		return logger
	}
	return middleware.LocalLogger{}
}

// Read a file of module versions that we should ignore because
// the are in the index but not stored in the proxy.
// Format of the file: each line is
//     module@version
func readProxyRemoved(ctx context.Context) {
	filename := config.GetEnv("GO_DISCOVERY_PROXY_REMOVED", "")
	if filename == "" {
		return
	}
	lines, err := readFileLines(filename)
	if err != nil {
		log.Fatal(ctx, err)
	}
	for _, line := range lines {
		etl.ProxyRemoved[line] = true
	}
	log.Infof(ctx, "read %d excluded module versions from %s", len(etl.ProxyRemoved), filename)
}

// populateExcluded adds each element of excludedPrefixes to the excluded_prefixes
// table if it isn't already present.
func populateExcluded(ctx context.Context, db *postgres.DB) {
	const excludedFilename = "excluded.txt"
	lines, err := readFileLines(excludedFilename)
	if err != nil {
		log.Fatal(ctx, err)
	}
	user := os.Getenv("USER")
	if user == "" {
		user = "etl"
	}
	for _, line := range lines {
		var prefix, reason string
		i := strings.IndexAny(line, " \t")
		if i >= 0 {
			prefix = line[:i]
			reason = strings.TrimSpace(line[i+1:])
		}
		if reason == "" {
			log.Fatalf(ctx, "missing reason in %s, line %q", excludedFilename, line)
		}
		fmt.Printf("#### %q ->\n", line)
		fmt.Printf(" prefix=%q, reason=%q\n", prefix, reason)
		present, err := db.IsExcluded(ctx, prefix)
		if err != nil {
			log.Fatalf(ctx, "db.IsExcluded(%q): %v", prefix, err)
		}
		if !present {
			if err := db.InsertExcludedPrefix(ctx, prefix, user, reason); err != nil {
				log.Fatalf(ctx, "db.InsertExcludedPrefix(%q, %q, %q): %v", prefix, user, reason, err)
			}
		}
	}
}

// readFileLines reads filename and returns its lines, trimmed of whitespace.
// Blank lines and lines whose first non-blank character is '#' are omitted.
func readFileLines(filename string) ([]string, error) {
	var lines []string
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	scan := bufio.NewScanner(f)
	for scan.Scan() {
		line := strings.TrimSpace(scan.Text())
		if line == "" || line[0] == '#' {
			continue
		}
		lines = append(lines, line)
	}
	if err := scan.Err(); err != nil {
		return nil, err
	}
	return lines, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"math/rand"
	"os"
	"sort"
	"time"

	"github.com/lucapette/fakedata/pkg/fakedata"
	flag "github.com/spf13/pflag"
)

var usage = `
  Usage: fakedata field1 [[field2] [field3]]
    [--tick d]
    [--max n]
    [--generators]
    [--format]

    fakedata --help

  Options:
    --generators    list all available generators
    --max n         generate data up to n [default: 10]
    --tick d        generate data every d [default: 10ms]
    --help          show help information
    --format f      generate data in f format [options: csv|tab, default: " "]
`

var generatorsFlag = flag.Bool("generators", false, "list all the generators")
var tickFlag = flag.Duration("tick", 10*time.Millisecond, "generate data every d milliseconds")
var maxFlag = flag.Int("max", 10, "generate up to n rows")
var helpFlag = flag.Bool("help", false, "print usage")
var formatFlag = flag.String("format", "", "Output format")

func main() {
	if *helpFlag {
		fmt.Printf(usage)
		os.Exit(0)
	}

	if *generatorsFlag {
		generators := fakedata.List()
		sort.Strings(generators)

		for _, name := range generators {
			fmt.Printf("%s\n", name)
		}
		os.Exit(0)
	}

	if len(flag.Args()) == 0 {
		fmt.Printf(usage)
		os.Exit(0)
	}

	rand.Seed(time.Now().UnixNano())

	tick := time.Tick(*tickFlag)

	total := 0

	for range tick {
		fmt.Print(fakedata.GenerateRow(flag.Args(), *formatFlag))

		if total++; total == *maxFlag {
			return
		}
	}
}

func init() {
	flag.Parse()
}
<commit_msg>Remove tick feature<commit_after>package main

import (
	"fmt"
	"math/rand"
	"os"
	"sort"
	"time"

	"github.com/lucapette/fakedata/pkg/fakedata"
	flag "github.com/spf13/pflag"
)

var usage = `
  Usage: fakedata field1 [[field2] [field3]]
    [--tick d]
    [--max n]
    [--generators]
    [--format]

    fakedata --help

  Options:
    --generators    list all available generators
    --max n         generate data up to n [default: 10]
    --help          show help information
    --format f      generate data in f format [options: csv|tab, default: " "]
`

var generatorsFlag = flag.Bool("generators", false, "list all the generators")
var maxFlag = flag.Int("max", 10, "generate up to n rows")
var helpFlag = flag.Bool("help", false, "print usage")
var formatFlag = flag.String("format", "", "Output format")

func main() {
	if *helpFlag {
		fmt.Printf(usage)
		os.Exit(0)
	}

	if *generatorsFlag {
		generators := fakedata.List()
		sort.Strings(generators)

		for _, name := range generators {
			fmt.Printf("%s\n", name)
		}
		os.Exit(0)
	}

	if len(flag.Args()) == 0 {
		fmt.Printf(usage)
		os.Exit(0)
	}

	rand.Seed(time.Now().UnixNano())

	for i := 0; i < *maxFlag; i++ {
		fmt.Print(fakedata.GenerateRow(flag.Args(), *formatFlag))
	}
}

func init() {
	flag.Parse()
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"net/http"
	"net/http/pprof"
	"net/url"
	"time"

	"github.com/Sirupsen/logrus"
	negronilogrus "github.com/meatballhat/negroni-logrus"
	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/urfave/negroni"
	"github.com/wyattjoh/ims/cmd/ims/routes"
	"github.com/wyattjoh/ims/internal/image/provider"
)

// GetProvider gets the image provider to use for the resize handler. If the
// origin is not provided, it defaults to the filesysytem provider with the
// specified directory.
func GetProvider(directory, origin string) (provider.Provider, error) {

	// By default, we'll try to use the directory resize, otherwise, if the origin
	// url is provided, use it.
	if origin == "" {
		logrus.WithField("directory", directory).Debug("serving from the filesystem")

		return provider.Filesystem{Dir: http.Dir(directory)}, nil
	}

	logrus.WithField("origin", origin).Debug("serving from the origin")

	originURL, err := url.Parse(origin)
	if err != nil {
		return nil, errors.Wrap(err, "can't parse the origin url")
	}

	return provider.Origin{URL: originURL}, nil

}

type ServerOpts struct {
	Addr           string
	Debug          bool
	DisableMetrics bool
	Directory      string
	Origin         string
	CacheTimeout   time.Duration
}

// Serve creates and starts a new server to provide image resizing services.
func Serve(opts *ServerOpts) error {
	mux := http.NewServeMux()

	if opts.CacheTimeout != 0 {
		logrus.WithField("timeout", opts.CacheTimeout.String()).Debug("cache headers enabled")
	} else {
		logrus.Debug("cache headers disabled")
	}

	// When debug mode is enabled, mount the debug handlers on this router.
	if opts.Debug {
		mux.Handle("/debug/pprof/", http.HandlerFunc(pprof.Index))
		mux.Handle("/debug/pprof/cmdline", http.HandlerFunc(pprof.Cmdline))
		mux.Handle("/debug/pprof/profile", http.HandlerFunc(pprof.Profile))
		mux.Handle("/debug/pprof/symbol", http.HandlerFunc(pprof.Symbol))
		mux.Handle("/debug/pprof/trace", http.HandlerFunc(pprof.Trace))
	}

	// Get the image provider.
	p, err := GetProvider(opts.Directory, opts.Origin)
	if err != nil {
		return err
	}

	if !opts.DisableMetrics {

		// Mount the resize handler on the mux with the instrumentation wrapped on
		// the handler.
		mux.HandleFunc("/", prometheus.InstrumentHandler("image", routes.Resize(opts.CacheTimeout, p)))

		// Register the prometheus metrics handler.
		mux.Handle("/metrics", prometheus.Handler())
	} else {

		// Mount the resize handler on the mux.
		mux.HandleFunc("/", routes.Resize(opts.CacheTimeout, p))
	}

	// Create the negroni middleware bundle.
	n := negroni.New(negroni.NewRecovery(), negronilogrus.NewMiddleware())

	// Attach the mux.
	n.UseHandler(mux)

	logrus.WithField("address", opts.Addr).Info("Now listening")
	return http.ListenAndServe(opts.Addr, n)
}

func main() {
	var (
		debug           = flag.Bool("debug", false, "enable debug logging and pprof routes")
		listenAddr      = flag.String("listen-addr", "0.0.0.0:8080", "the address to listen for new connections on")
		imagesDirectory = flag.String("images-dir", "images", "the location on the filesystem to load images from")
		originURL       = flag.String("origin-url", "", "url for the origin server to pull images from")
		timeout         = flag.Duration("timeout", 15*time.Minute, "used to set the cache control max age headers, set to 0 to disable")
		disableMetrics  = flag.Bool("disable-metrics", false, "disable the prometheus metrics")
	)

	flag.Parse()

	if *debug {
		logrus.SetLevel(logrus.DebugLevel)
	}

	// Setup the server options.
	opts := &ServerOpts{
		Addr:           *listenAddr,
		Debug:          *debug,
		DisableMetrics: *disableMetrics,
		Directory:      *imagesDirectory,
		Origin:         *originURL,
		CacheTimeout:   *timeout,
	}

	if err := Serve(opts); err != nil {
		logrus.Fatalf("Could not serve: %s", err)
	}
}
<commit_msg>linting<commit_after>package main

import (
	"flag"
	"net/http"
	"net/http/pprof"
	"net/url"
	"time"

	"github.com/Sirupsen/logrus"
	negronilogrus "github.com/meatballhat/negroni-logrus"
	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/urfave/negroni"
	"github.com/wyattjoh/ims/cmd/ims/routes"
	"github.com/wyattjoh/ims/internal/image/provider"
)

// GetProvider gets the image provider to use for the resize handler. If the
// origin is not provided, it defaults to the filesysytem provider with the
// specified directory.
func GetProvider(directory, origin string) (provider.Provider, error) {

	// By default, we'll try to use the directory resize, otherwise, if the origin
	// url is provided, use it.
	if origin == "" {
		logrus.WithField("directory", directory).Debug("serving from the filesystem")

		return provider.Filesystem{Dir: http.Dir(directory)}, nil
	}

	logrus.WithField("origin", origin).Debug("serving from the origin")

	originURL, err := url.Parse(origin)
	if err != nil {
		return nil, errors.Wrap(err, "can't parse the origin url")
	}

	return provider.Origin{URL: originURL}, nil

}

// ServerOpts is the options for starting a new Server,
type ServerOpts struct {

	// Addr is the address to listen for http requests on.
	Addr string

	// Debug enables pprof endpoints and debug logs.
	Debug bool

	// DisableMetrics disables Prometheus endpoints.
	DisableMetrics bool

	// Directory is the folder in which images are served out of.
	Directory string

	// Origin is the url that is the base url for images and will act as the
	// provider.
	Origin string

	// CacheTimeout is the time that images will have cache headers for when
	// writing them out to the http response.
	CacheTimeout time.Duration
}

// Serve creates and starts a new server to provide image resizing services.
func Serve(opts *ServerOpts) error {
	mux := http.NewServeMux()

	if opts.CacheTimeout != 0 {
		logrus.WithField("timeout", opts.CacheTimeout.String()).Debug("cache headers enabled")
	} else {
		logrus.Debug("cache headers disabled")
	}

	// When debug mode is enabled, mount the debug handlers on this router.
	if opts.Debug {
		mux.Handle("/debug/pprof/", http.HandlerFunc(pprof.Index))
		mux.Handle("/debug/pprof/cmdline", http.HandlerFunc(pprof.Cmdline))
		mux.Handle("/debug/pprof/profile", http.HandlerFunc(pprof.Profile))
		mux.Handle("/debug/pprof/symbol", http.HandlerFunc(pprof.Symbol))
		mux.Handle("/debug/pprof/trace", http.HandlerFunc(pprof.Trace))
	}

	// Get the image provider.
	p, err := GetProvider(opts.Directory, opts.Origin)
	if err != nil {
		return err
	}

	if !opts.DisableMetrics {

		// Mount the resize handler on the mux with the instrumentation wrapped on
		// the handler.
		mux.HandleFunc("/", prometheus.InstrumentHandler("image", routes.Resize(opts.CacheTimeout, p)))

		// Register the prometheus metrics handler.
		mux.Handle("/metrics", prometheus.Handler())
	} else {

		// Mount the resize handler on the mux.
		mux.HandleFunc("/", routes.Resize(opts.CacheTimeout, p))
	}

	// Create the negroni middleware bundle.
	n := negroni.New(negroni.NewRecovery(), negronilogrus.NewMiddleware())

	// Attach the mux.
	n.UseHandler(mux)

	logrus.WithField("address", opts.Addr).Info("Now listening")
	return http.ListenAndServe(opts.Addr, n)
}

func main() {
	var (
		debug           = flag.Bool("debug", false, "enable debug logging and pprof routes")
		listenAddr      = flag.String("listen-addr", "0.0.0.0:8080", "the address to listen for new connections on")
		imagesDirectory = flag.String("images-dir", "images", "the location on the filesystem to load images from")
		originURL       = flag.String("origin-url", "", "url for the origin server to pull images from")
		timeout         = flag.Duration("timeout", 15*time.Minute, "used to set the cache control max age headers, set to 0 to disable")
		disableMetrics  = flag.Bool("disable-metrics", false, "disable the prometheus metrics")
	)

	flag.Parse()

	if *debug {
		logrus.SetLevel(logrus.DebugLevel)
	}

	// Setup the server options.
	opts := &ServerOpts{
		Addr:           *listenAddr,
		Debug:          *debug,
		DisableMetrics: *disableMetrics,
		Directory:      *imagesDirectory,
		Origin:         *originURL,
		CacheTimeout:   *timeout,
	}

	if err := Serve(opts); err != nil {
		logrus.Fatalf("Could not serve: %s", err)
	}
}
<|endoftext|>