<commit_before>package routernode

import (
	"fmt"
	"math/rand"
	"net/http"
	"sync/atomic"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/julienschmidt/httprouter"

	"github.com/jacksontj/dataman/src/query"
	"github.com/jacksontj/dataman/src/router_node/metadata"
)

// This node is responsible for routing requests to the appropriate storage node
// This is also responsible for maintaining schema, indexes, etc. from the metadata store
type RouterNode struct {
	Config    *Config
	MetaStore *MetadataStore

	meta atomic.Value

	// background sync stuff
	stop chan struct{}
	Sync chan struct{}
}

func NewRouterNode(config *Config) (*RouterNode, error) {
	metaStore, err := NewMetadataStore(config)
	if err != nil {
		return nil, err
	}
	node := &RouterNode{
		Config:    config,
		MetaStore: metaStore,
	}

	// TODO: check that it worked?
	// Before returning we should get the metadata from the metadata store
	node.FetchMeta()
	go node.background()

	// TODO: background goroutine to re-fetch every interval (with some mechanism to trigger on-demand)

	return node, nil
}

// TODO: have a stop?
func (s *RouterNode) Start() error {
	// initialize the http api (since at this point we are ready to go!
	router := httprouter.New()
	api := NewHTTPApi(s)
	api.Start(router)

	return http.ListenAndServe(s.Config.HTTP.Addr, router)
}

func (s *RouterNode) GetMeta() *metadata.Meta {
	return s.meta.Load().(*metadata.Meta)
}

func (s *RouterNode) background() {
	interval := time.Second // TODO: configurable interval
	ticker := time.NewTicker(interval)

	for {
		select {
		case <-ticker.C: // time based trigger, in case of error etc.
			s.FetchMeta()
		case <-s.Sync: // event based trigger, so we can get stuff to disk ASAP
			s.FetchMeta()
			// since we where just triggered, lets reset the interval
			ticker = time.NewTicker(interval)
		}
	}
}

// This method will create a new `Databases` map and swap it in
func (s *RouterNode) FetchMeta() error {
	// First we need to determine all the databases that we are responsible for
	// TODO: lots of error handling required

	// TODO: support errors
	meta := s.MetaStore.GetMeta()
	if meta != nil {
		s.meta.Store(meta)
	}
	logrus.Debugf("Loaded meta: %v", meta)
	return nil
}

func (s *RouterNode) HandleQueries(queries []map[query.QueryType]query.QueryArgs) []*query.Result {
	// TODO: we should actually do these in parallel (potentially with some
	// config of *how* parallel)
	results := make([]*query.Result, len(queries))

	// We specifically want to load this once for the batch so we don't have mixed
	// schema information across this batch of queries
	meta := s.GetMeta()

	for i, queryMap := range queries {
		// We only allow a single method to be defined per item
		if len(queryMap) == 1 {
			for queryType, queryArgs := range queryMap {
				// Switch between read and write operations
				switch queryType {
				// Write operations
				case query.Set:
					fallthrough
				case query.Insert:
					fallthrough
				case query.Update:
					fallthrough
				case query.Delete:
					results[i] = s.handleWrite(meta, queryType, queryArgs)

				// Read operations
				case query.Get:
					fallthrough
				case query.Filter:
					results[i] = s.handleRead(meta, queryType, queryArgs)

					// All other operations should error
				default:
					results[i] = &query.Result{Error: "Unkown query type: " + string(queryType)}
				}
			}

		} else {
			results[i] = &query.Result{
				Error: fmt.Sprintf("Only one QueryType supported per query: %v -- %v", queryMap, queries),
			}
		}
	}
	return results
}

func (s *RouterNode) handleRead(meta *metadata.Meta, queryType query.QueryType, queryArgs query.QueryArgs) *query.Result {
	database, ok := meta.Databases[queryArgs["db"].(string)]
	if !ok {
		return &query.Result{Error: "Unknown db " + queryArgs["db"].(string)}
	}
	collection, ok := database.Collections[queryArgs["collection"].(string)]
	if !ok {
		return &query.Result{Error: "Unknown collection " + queryArgs["collection"].(string)}
	}

	// TODO: get collection? Later we'll want to do shard keys which aren't "_id"
	// and to do that we'll need the collection metadata

	// Once we have the metadata all found we need to do the following:
	//      - Authentication/authorization
	//      - Cache
	//      - Sharding
	//      - Replicas

	//TODO:Authentication/authorization
	//TODO:Cache (configurable)

	// Sharding
	var shards []*metadata.DatastoreShard
	switch queryType {
	case query.Get:
		partition := collection.GetPartition(int64(queryArgs["_id"].(float64)))
		// TODO: configurable shard key
		shardNum, err := partition.ShardFunc(queryArgs[partition.ShardConfig.Key], len(database.Datastore.Shards))
		if err != nil {
			return &query.Result{Error: err.Error()}
		}
		shards = []*metadata.DatastoreShard{database.Datastore.Shards[shardNum]}
	case query.Filter:
		// if there is only one partition and we have our shard key, we can be more specific
		if shardKeyValue, ok := queryArgs["filter"].(map[string]interface{})[collection.Partitions[0].ShardConfig.Key]; len(collection.Partitions) == 1 && ok {
			shardNum, err := collection.Partitions[0].ShardFunc(shardKeyValue, len(database.Datastore.Shards))
			if err != nil {
				return &query.Result{Error: err.Error()}
			}
			shards = []*metadata.DatastoreShard{database.Datastore.Shards[shardNum]}
		} else {
			shards = database.Datastore.Shards
		}
	}

	shardResults := make([]*query.Result, len(shards))

	// TODO: parallel
	for i, shard := range shards {
		// TODO: replicas -- add args for slave etc.
		if result, err := QuerySingle(shard.Replicas.GetMaster().Store, &query.Query{queryType, queryArgs}); err == nil {
			shardResults[i] = result
		} else {
			shardResults[i] = &query.Result{Error: err.Error()}
		}

	}

	return query.MergeResult(shardResults...)
}

func (s *RouterNode) handleWrite(meta *metadata.Meta, queryType query.QueryType, queryArgs query.QueryArgs) *query.Result {
	database, ok := meta.Databases[queryArgs["db"].(string)]
	if !ok {
		return &query.Result{Error: "Unknown db " + queryArgs["db"].(string)}
	}
	collection, ok := database.Collections[queryArgs["collection"].(string)]
	if !ok {
		return &query.Result{Error: "Unknown collection " + queryArgs["collection"].(string)}
	}

	// TODO: get collection? Later we'll want to do shard keys which aren't "_id"
	// and to do that we'll need the collection metadata

	// Once we have the metadata all found we need to do the following:
	//      - Authentication/authorization
	//      - Cache
	//      - Sharding

	// TODO: Authentication/authorization
	// TODO: Cache poison

	// Sharding

	// TODO: eventually we'll want to be more sophisticated and do this same thing if there
	// are a set of id's we can derive from the original query, so we can do a limited
	// scatter-gather. For now we'll either know the specific shard, or not (for ease of implementation)

	switch queryType {
	// Write operations
	case query.Set:
		// If there is an "_id" present, then this is just a very specific update -- so we can find our specific shard
		if id, ok := queryArgs["record"].(map[string]interface{})["_id"]; ok {
			partition := collection.GetPartition(int64(id.(float64)))
			// TODO: configurable shard key
			shardNum, err := partition.ShardFunc(queryArgs[partition.ShardConfig.Key], len(database.Datastore.Shards))
			if err != nil {
				return &query.Result{Error: err.Error()}
			}

			// TODO: replicas -- add args for slave etc.
			if result, err := QuerySingle(database.Datastore.Shards[shardNum].Replicas.GetMaster().Store, &query.Query{queryType, queryArgs}); err == nil {
				return result
			} else {
				return &query.Result{Error: err.Error()}
			}
		} else { // Otherwise this is actually an insert, so we'll let it fall through to be handled as such
			// TODO: what do we want to do for brand new things?
			// TODO: consolidate into a single insert method
			// We want to RR between the shards for new inserts
			// TODO: counter instead of random (have to deal with the meta being syncd all over)
			shardNum := rand.Int63n(int64(len(database.Datastore.Shards)))

			// TODO: replicas -- add args for slave etc.
			result, err := QuerySingle(
				database.Datastore.Shards[shardNum].Replicas.GetMaster().Store,
				&query.Query{queryType, queryArgs},
			)

			if err == nil {
				return result
			} else {
				return &query.Result{Error: err.Error()}
			}
		}
	// TODO: what do we want to do for brand new things?
	case query.Insert:
		// TODO: not a good assumption if the key is _id (for example) since we won't know
		// Since this is a new item, we'll fall into the newest partition (presumably?)
		partition := collection.Partitions[len(collection.Partitions)-1]

		var shardNum int
		// If the key is _id (our monotomically increasing number) then we'll just RR between the shards
		if partition.ShardConfig.Key == "_id" {
			// TODO: consolidate into a single insert method
			// We want to RR between the shards for new inserts
			// TODO: counter instead of random (have to deal with the meta being syncd all over)
			shardNum = rand.Intn(len(database.Datastore.Shards))
		} else { // Otherwise we *must* have the shard key defined to do the sharding
			shardKeyValue, ok := queryArgs["record"].(map[string]interface{})[partition.ShardConfig.Key]
			if !ok {
				return &query.Result{Error: "Record missing shardKey " + partition.ShardConfig.Key}
			}
			var err error
			shardNum, err = partition.ShardFunc(shardKeyValue, len(database.Datastore.Shards))
			if err != nil {
				return &query.Result{Error: err.Error()}
			}
		}

		result, err := QuerySingle(
			// TODO: replicas -- add args for slave etc.
			database.Datastore.Shards[shardNum].Replicas.GetMaster().Store,
			&query.Query{queryType, queryArgs},
		)

		if err == nil {
			return result
		} else {
			return &query.Result{Error: err.Error()}
		}
	case query.Update:
		// If there is an "_id"_ defined, then we can send this to a single shard
		if id, ok := queryArgs["filter"].(map[string]interface{})["_id"]; ok {
			partition := collection.GetPartition(int64(id.(float64)))
			// TODO: configurable shard key
			shardNum, err := partition.ShardFunc(queryArgs[partition.ShardConfig.Key], len(database.Datastore.Shards))
			if err != nil {
				return &query.Result{Error: err.Error()}
			}
			// TODO: replicas -- add args for slave etc.
			if result, err := QuerySingle(database.Datastore.Shards[shardNum].Replicas.GetMaster().Store, &query.Query{queryType, queryArgs}); err == nil {
				return result
			} else {
				return &query.Result{Error: err.Error()}
			}

		} else { // Otherwise we need to send this query to all shards to let them handle it
			shardResults := make([]*query.Result, len(database.Datastore.Shards))

			// TODO: parallel
			for i, shard := range database.Datastore.Shards {
				// TODO: replicas -- add args for slave etc.
				if result, err := QuerySingle(shard.Replicas.GetMaster().Store, &query.Query{queryType, queryArgs}); err == nil {
					shardResults[i] = result
				} else {
					shardResults[i] = &query.Result{Error: err.Error()}
				}

			}

			return query.MergeResult(shardResults...)
		}
	case query.Delete:
		partition := collection.GetPartition(int64(queryArgs["_id"].(float64)))
		shardNum, err := partition.ShardFunc(queryArgs[partition.ShardConfig.Key], len(database.Datastore.Shards))
		if err != nil {
			return &query.Result{Error: err.Error()}
		}
		// TODO: replicas -- add args for slave etc.
		if result, err := QuerySingle(database.Datastore.Shards[shardNum].Replicas.GetMaster().Store, &query.Query{queryType, queryArgs}); err == nil {
			return result
		} else {
			return &query.Result{Error: err.Error()}
		}

	}

	return nil
}
<commit_msg>More sharding work<commit_after>package routernode

import (
	"fmt"
	"math/rand"
	"net/http"
	"sync/atomic"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/julienschmidt/httprouter"

	"github.com/jacksontj/dataman/src/query"
	"github.com/jacksontj/dataman/src/router_node/metadata"
)

// This node is responsible for routing requests to the appropriate storage node
// This is also responsible for maintaining schema, indexes, etc. from the metadata store
type RouterNode struct {
	Config    *Config
	MetaStore *MetadataStore

	meta atomic.Value

	// background sync stuff
	stop chan struct{}
	Sync chan struct{}
}

func NewRouterNode(config *Config) (*RouterNode, error) {
	metaStore, err := NewMetadataStore(config)
	if err != nil {
		return nil, err
	}
	node := &RouterNode{
		Config:    config,
		MetaStore: metaStore,
	}

	// TODO: check that it worked?
	// Before returning we should get the metadata from the metadata store
	node.FetchMeta()
	go node.background()

	// TODO: background goroutine to re-fetch every interval (with some mechanism to trigger on-demand)

	return node, nil
}

// TODO: have a stop?
func (s *RouterNode) Start() error {
	// initialize the http api (since at this point we are ready to go!
	router := httprouter.New()
	api := NewHTTPApi(s)
	api.Start(router)

	return http.ListenAndServe(s.Config.HTTP.Addr, router)
}

func (s *RouterNode) GetMeta() *metadata.Meta {
	return s.meta.Load().(*metadata.Meta)
}

func (s *RouterNode) background() {
	interval := time.Second // TODO: configurable interval
	ticker := time.NewTicker(interval)

	for {
		select {
		case <-ticker.C: // time based trigger, in case of error etc.
			s.FetchMeta()
		case <-s.Sync: // event based trigger, so we can get stuff to disk ASAP
			s.FetchMeta()
			// since we where just triggered, lets reset the interval
			ticker = time.NewTicker(interval)
		}
	}
}

// This method will create a new `Databases` map and swap it in
func (s *RouterNode) FetchMeta() error {
	// First we need to determine all the databases that we are responsible for
	// TODO: lots of error handling required

	// TODO: support errors
	meta := s.MetaStore.GetMeta()
	if meta != nil {
		s.meta.Store(meta)
	}
	logrus.Debugf("Loaded meta: %v", meta)
	return nil
}

func (s *RouterNode) HandleQueries(queries []map[query.QueryType]query.QueryArgs) []*query.Result {
	// TODO: we should actually do these in parallel (potentially with some
	// config of *how* parallel)
	results := make([]*query.Result, len(queries))

	// We specifically want to load this once for the batch so we don't have mixed
	// schema information across this batch of queries
	meta := s.GetMeta()

	for i, queryMap := range queries {
		// We only allow a single method to be defined per item
		if len(queryMap) == 1 {
			for queryType, queryArgs := range queryMap {
				// Switch between read and write operations
				switch queryType {
				// Write operations
				case query.Set:
					fallthrough
				case query.Insert:
					fallthrough
				case query.Update:
					fallthrough
				case query.Delete:
					results[i] = s.handleWrite(meta, queryType, queryArgs)

				// Read operations
				case query.Get:
					fallthrough
				case query.Filter:
					results[i] = s.handleRead(meta, queryType, queryArgs)

					// All other operations should error
				default:
					results[i] = &query.Result{Error: "Unkown query type: " + string(queryType)}
				}
			}

		} else {
			results[i] = &query.Result{
				Error: fmt.Sprintf("Only one QueryType supported per query: %v -- %v", queryMap, queries),
			}
		}
	}
	return results
}

func (s *RouterNode) handleRead(meta *metadata.Meta, queryType query.QueryType, queryArgs query.QueryArgs) *query.Result {
	database, ok := meta.Databases[queryArgs["db"].(string)]
	if !ok {
		return &query.Result{Error: "Unknown db " + queryArgs["db"].(string)}
	}
	collection, ok := database.Collections[queryArgs["collection"].(string)]
	if !ok {
		return &query.Result{Error: "Unknown collection " + queryArgs["collection"].(string)}
	}

	// TODO: get collection? Later we'll want to do shard keys which aren't "_id"
	// and to do that we'll need the collection metadata

	// Once we have the metadata all found we need to do the following:
	//      - Authentication/authorization
	//      - Cache
	//      - Sharding
	//      - Replicas

	//TODO:Authentication/authorization
	//TODO:Cache (configurable)

	// Sharding
	var shards []*metadata.DatastoreShard
	switch queryType {
	case query.Get:
		partition := collection.GetPartition(int64(queryArgs["_id"].(float64)))
		// TODO: configurable shard key
		shardNum, err := partition.ShardFunc(queryArgs[partition.ShardConfig.Key], len(database.Datastore.Shards))
		if err != nil {
			return &query.Result{Error: err.Error()}
		}
		shards = []*metadata.DatastoreShard{database.Datastore.Shards[shardNum]}
	case query.Filter:
		// if there is only one partition and we have our shard key, we can be more specific
		if shardKeyValue, ok := queryArgs["filter"].(map[string]interface{})[collection.Partitions[0].ShardConfig.Key]; len(collection.Partitions) == 1 && ok {
			shardNum, err := collection.Partitions[0].ShardFunc(shardKeyValue, len(database.Datastore.Shards))
			if err != nil {
				return &query.Result{Error: err.Error()}
			}
			shards = []*metadata.DatastoreShard{database.Datastore.Shards[shardNum]}
			fmt.Printf("here %v %v %d\n", queryArgs, shards[0], shardNum)
		} else {
			shards = database.Datastore.Shards
		}
	}

	shardResults := make([]*query.Result, len(shards))

	// TODO: parallel
	for i, shard := range shards {
		// TODO: replicas -- add args for slave etc.
		if result, err := QuerySingle(shard.Replicas.GetMaster().Store, &query.Query{queryType, queryArgs}); err == nil {
			shardResults[i] = result
		} else {
			shardResults[i] = &query.Result{Error: err.Error()}
		}

	}

	return query.MergeResult(shardResults...)
}

func (s *RouterNode) handleWrite(meta *metadata.Meta, queryType query.QueryType, queryArgs query.QueryArgs) *query.Result {
	database, ok := meta.Databases[queryArgs["db"].(string)]
	if !ok {
		return &query.Result{Error: "Unknown db " + queryArgs["db"].(string)}
	}
	collection, ok := database.Collections[queryArgs["collection"].(string)]
	if !ok {
		return &query.Result{Error: "Unknown collection " + queryArgs["collection"].(string)}
	}

	// TODO: get collection? Later we'll want to do shard keys which aren't "_id"
	// and to do that we'll need the collection metadata

	// Once we have the metadata all found we need to do the following:
	//      - Authentication/authorization
	//      - Cache
	//      - Sharding

	// TODO: Authentication/authorization
	// TODO: Cache poison

	// Sharding

	// TODO: eventually we'll want to be more sophisticated and do this same thing if there
	// are a set of id's we can derive from the original query, so we can do a limited
	// scatter-gather. For now we'll either know the specific shard, or not (for ease of implementation)

	switch queryType {
	// Write operations
	case query.Set:
		// If there is an "_id" present, then this is just a very specific update -- so we can find our specific shard
		if id, ok := queryArgs["record"].(map[string]interface{})["_id"]; ok {
			partition := collection.GetPartition(int64(id.(float64)))
			// TODO: configurable shard key
			shardNum, err := partition.ShardFunc(queryArgs[partition.ShardConfig.Key], len(database.Datastore.Shards))
			if err != nil {
				return &query.Result{Error: err.Error()}
			}

			// TODO: replicas -- add args for slave etc.
			if result, err := QuerySingle(database.Datastore.Shards[shardNum].Replicas.GetMaster().Store, &query.Query{queryType, queryArgs}); err == nil {
				return result
			} else {
				return &query.Result{Error: err.Error()}
			}
		} else { // Otherwise this is actually an insert, so we'll let it fall through to be handled as such
			// TODO: what do we want to do for brand new things?
			// TODO: consolidate into a single insert method
			// We assume that it is always safe to do the insert on the "newest"
			// partition. This should be safe as the storage_node is responsible for
			// generating the "_id" which will be within a partition range that the node
			// was responsible for. Specifically this means that even if we have a newer partition
			// than the storage_node knows of, it will do the insert with an _id that falls into
			// a partition which is owned by that storage_node.
			partition := collection.Partitions[len(collection.Partitions)-1]

			var shardNum int
			// If the key is _id (our monotomically increasing number) then we'll just RR between the shards
			if partition.ShardConfig.Key == "_id" {
				// TODO: consolidate into a single insert method
				// We want to RR between the shards for new inserts
				// TODO: counter instead of random (have to deal with the meta being syncd all over)
				shardNum = rand.Intn(len(database.Datastore.Shards))
			} else { // Otherwise we *must* have the shard key defined to do the sharding
				shardKeyValue, ok := queryArgs["record"].(map[string]interface{})[partition.ShardConfig.Key]
				if !ok {
					return &query.Result{Error: "Record missing shardKey " + partition.ShardConfig.Key}
				}
				var err error
				shardNum, err = partition.ShardFunc(shardKeyValue, len(database.Datastore.Shards))
				if err != nil {
					return &query.Result{Error: err.Error()}
				}
			}

			result, err := QuerySingle(
				// TODO: replicas -- add args for slave etc.
				database.Datastore.Shards[shardNum].Replicas.GetMaster().Store,
				&query.Query{queryType, queryArgs},
			)

			if err == nil {
				return result
			} else {
				return &query.Result{Error: err.Error()}
			}
		}
	// TODO: what do we want to do for brand new things?
	case query.Insert:
		// We assume that it is always safe to do the insert on the "newest"
		// partition. This should be safe as the storage_node is responsible for
		// generating the "_id" which will be within a partition range that the node
		// was responsible for. Specifically this means that even if we have a newer partition
		// than the storage_node knows of, it will do the insert with an _id that falls into
		// a partition which is owned by that storage_node.
		partition := collection.Partitions[len(collection.Partitions)-1]

		var shardNum int
		// If the key is _id (our monotomically increasing number) then we'll just RR between the shards
		if partition.ShardConfig.Key == "_id" {
			// TODO: consolidate into a single insert method
			// We want to RR between the shards for new inserts
			// TODO: counter instead of random (have to deal with the meta being syncd all over)
			shardNum = rand.Intn(len(database.Datastore.Shards))
		} else { // Otherwise we *must* have the shard key defined to do the sharding
			shardKeyValue, ok := queryArgs["record"].(map[string]interface{})[partition.ShardConfig.Key]
			if !ok {
				return &query.Result{Error: "Record missing shardKey " + partition.ShardConfig.Key}
			}
			var err error
			shardNum, err = partition.ShardFunc(shardKeyValue, len(database.Datastore.Shards))
			if err != nil {
				return &query.Result{Error: err.Error()}
			}
		}

		result, err := QuerySingle(
			// TODO: replicas -- add args for slave etc.
			database.Datastore.Shards[shardNum].Replicas.GetMaster().Store,
			&query.Query{queryType, queryArgs},
		)

		if err == nil {
			return result
		} else {
			return &query.Result{Error: err.Error()}
		}
	case query.Update:
		// If there is an "_id"_ defined, then we can send this to a single shard
		if id, ok := queryArgs["filter"].(map[string]interface{})["_id"]; ok {
			partition := collection.GetPartition(int64(id.(float64)))
			// TODO: configurable shard key
			shardNum, err := partition.ShardFunc(queryArgs[partition.ShardConfig.Key], len(database.Datastore.Shards))
			if err != nil {
				return &query.Result{Error: err.Error()}
			}
			// TODO: replicas -- add args for slave etc.
			if result, err := QuerySingle(database.Datastore.Shards[shardNum].Replicas.GetMaster().Store, &query.Query{queryType, queryArgs}); err == nil {
				return result
			} else {
				return &query.Result{Error: err.Error()}
			}

		} else { // Otherwise we need to send this query to all shards to let them handle it
			shardResults := make([]*query.Result, len(database.Datastore.Shards))

			// TODO: parallel
			for i, shard := range database.Datastore.Shards {
				// TODO: replicas -- add args for slave etc.
				if result, err := QuerySingle(shard.Replicas.GetMaster().Store, &query.Query{queryType, queryArgs}); err == nil {
					shardResults[i] = result
				} else {
					shardResults[i] = &query.Result{Error: err.Error()}
				}

			}

			return query.MergeResult(shardResults...)
		}
	case query.Delete:
		partition := collection.GetPartition(int64(queryArgs["_id"].(float64)))
		shardNum, err := partition.ShardFunc(queryArgs[partition.ShardConfig.Key], len(database.Datastore.Shards))
		if err != nil {
			return &query.Result{Error: err.Error()}
		}
		// TODO: replicas -- add args for slave etc.
		if result, err := QuerySingle(database.Datastore.Shards[shardNum].Replicas.GetMaster().Store, &query.Query{queryType, queryArgs}); err == nil {
			return result
		} else {
			return &query.Result{Error: err.Error()}
		}

	}

	return nil
}
<|endoftext|>


<commit_before>package utee

import (
	"errors"
	"github.com/garyburd/redigo/redis"
	"log"
	"time"
)

var ErrFull = errors.New("queue is full")

type MemQueue chan interface{}

func NewMemQueue(cap int) MemQueue {
	return make(chan interface{}, cap)
}

func NewLeakMemQueue(cap, concurrent int, cb func(interface{})) MemQueue {
	q := NewMemQueue(cap)

	f := func() {
		for {
			cb(q.Deq())
		}
	}

	for i := 0; i < concurrent; i++ {
		go f()
	}
	return q
}

//enqueue, block if queue is full
func (p MemQueue) EnqBlocking(data interface{}) {
	p <- data
}

//enqueue, return error if queue is full
func (p MemQueue) Enq(data interface{}) error {
	select {
	case p <- data:
	default:
		return ErrFull
	}
	return nil
}

func (p MemQueue) Deq() interface{} {
	return <-p
}

//dequeue less than n in a batch
func (p MemQueue) DeqN(n int) []interface{} {
	if n <= 0 {
		log.Println("[MemQueue] deqn err, n must > 0")
		return nil
	}

	var l []interface{}

	select {
	case data := <-p:
		l = append(l, data)
		if len(l) == n {
			break
		}
	default:
		break
	}

	return l
}

func (p MemQueue) Len() int {
	return len(p)
}

func (p MemQueue) Cap() int {
	return cap(p)
}

func qname(name string) string {
	return "q" + name
}

type SimpleRedisQueue struct {
	name   string
	pool   *redis.Pool
	buffer MemQueue
	batch  int
}

//redis queue with optional memory buffer
//server: redis server address
//auth: redis auth
//name: queue name in redis
//concurrent: concurrent number redis enqueue operation, must >=1
//enqBatch: batch enqueue number, must >=1
//buffer: memory buffer capacity, must >= 0
func NewSimpleRedisQueue(server, auth, name string, concurrent, enqBatch, buffer int) *SimpleRedisQueue {
	if concurrent < 1 {
		log.Fatal("concurrent must >= 1")
	}
	if enqBatch < 1 {
		log.Fatal("batch must >= 1")
	}

	if buffer < 0 {
		log.Fatal("buffer must >= 0")
	}

	q := &SimpleRedisQueue{
		name:   qname(name),
		pool:   CreateRedisPool(concurrent, server, auth),
		buffer: NewMemQueue(buffer),
		batch:  enqBatch,
	}
	for i := 0; i < concurrent; i++ {
		go q.enqLoop()
	}
	return q
}

func (p *SimpleRedisQueue) enqLoop() {
	for {
		l := p.buffer.DeqN(p.batch)
		if len(l) > 0 {
			if err := p.enqBatch(l); err != nil {
				log.Println("[SimpleRedisQueue enqLoop] err ", err)
			}
		} else {
			time.Sleep(100 * time.Millisecond)
		}
	}
}

func (p *SimpleRedisQueue) enqBatch(l []interface{}) error {
	c := p.pool.Get()
	defer c.Close()
	for _, data := range l {
		if err := c.Send("RPUSH", p.name, data); err != nil {
			log.Println("[SimpleRedisQueue enqBatch] err :", err)
		}
	}
	return c.Flush()
}

func (p *SimpleRedisQueue) Len() (int, error) {
	c := p.pool.Get()
	defer c.Close()

	i, err := redis.Int(c.Do("LLEN", p.name))

	if err != nil && err.Error() == "redigo: nil returned" {
		//expire
		return 0, nil
	}
	return i, err
}

//enqueue, block if buffer is full
func (p *SimpleRedisQueue) EnqBlocking(data interface{}) {
	p.buffer.EnqBlocking(data)
}

//enqueue, return error if buffer is full
func (p *SimpleRedisQueue) Enq(data interface{}) error {
	return p.buffer.Enq(data)
}

func (p *SimpleRedisQueue) Deq() (interface{}, error) {
	c := p.pool.Get()
	defer c.Close()
	return c.Do("LPOP", p.name)
}

func (p *SimpleRedisQueue) BufferLen() int {
	return p.buffer.Len()
}

func (p *SimpleRedisQueue) BufferCap() int {
	return p.buffer.Cap()
}

func (p *SimpleRedisQueue) PollSize() int {
	return p.pool.ActiveCount()
}
<commit_msg>rename<commit_after>package utee

import (
	"errors"
	"github.com/garyburd/redigo/redis"
	"log"
	"time"
)

var ErrFull = errors.New("queue is full")

type MemQueue chan interface{}

func NewMemQueue(cap int) MemQueue {
	return make(chan interface{}, cap)
}

//create memory queue, auto-leak element concurrently to worker
func NewLeakMemQueue(cap, concurrent int, worker func(interface{})) MemQueue {
	q := NewMemQueue(cap)

	f := func() {
		for {
			worker(q.Deq())
		}
	}

	for i := 0; i < concurrent; i++ {
		go f()
	}
	return q
}

//enqueue, block if queue is full
func (p MemQueue) EnqBlocking(data interface{}) {
	p <- data
}

//enqueue, return error if queue is full
func (p MemQueue) Enq(data interface{}) error {
	select {
	case p <- data:
	default:
		return ErrFull
	}
	return nil
}

func (p MemQueue) Deq() interface{} {
	return <-p
}

//dequeue less than n in a batch
func (p MemQueue) DeqN(n int) []interface{} {
	if n <= 0 {
		log.Println("[MemQueue] deqn err, n must > 0")
		return nil
	}

	var l []interface{}

	select {
	case data := <-p:
		l = append(l, data)
		if len(l) == n {
			break
		}
	default:
		break
	}

	return l
}

func (p MemQueue) Len() int {
	return len(p)
}

func (p MemQueue) Cap() int {
	return cap(p)
}

func qname(name string) string {
	return "q" + name
}

type SimpleRedisQueue struct {
	name   string
	pool   *redis.Pool
	buffer MemQueue
	batch  int
}

//redis queue with optional memory buffer
//server: redis server address
//auth: redis auth
//name: queue name in redis
//concurrent: concurrent number redis enqueue operation, must >=1
//enqBatch: batch enqueue number, must >=1
//buffer: memory buffer capacity, must >= 0
func NewSimpleRedisQueue(server, auth, name string, concurrent, enqBatch, buffer int) *SimpleRedisQueue {
	if concurrent < 1 {
		log.Fatal("concurrent must >= 1")
	}
	if enqBatch < 1 {
		log.Fatal("batch must >= 1")
	}

	if buffer < 0 {
		log.Fatal("buffer must >= 0")
	}

	q := &SimpleRedisQueue{
		name:   qname(name),
		pool:   CreateRedisPool(concurrent, server, auth),
		buffer: NewMemQueue(buffer),
		batch:  enqBatch,
	}
	for i := 0; i < concurrent; i++ {
		go q.enqLoop()
	}
	return q
}

func (p *SimpleRedisQueue) enqLoop() {
	for {
		l := p.buffer.DeqN(p.batch)
		if len(l) > 0 {
			if err := p.enqBatch(l); err != nil {
				log.Println("[SimpleRedisQueue enqLoop] err ", err)
			}
		} else {
			time.Sleep(100 * time.Millisecond)
		}
	}
}

func (p *SimpleRedisQueue) enqBatch(l []interface{}) error {
	c := p.pool.Get()
	defer c.Close()
	for _, data := range l {
		if err := c.Send("RPUSH", p.name, data); err != nil {
			log.Println("[SimpleRedisQueue enqBatch] err :", err)
		}
	}
	return c.Flush()
}

func (p *SimpleRedisQueue) Len() (int, error) {
	c := p.pool.Get()
	defer c.Close()

	i, err := redis.Int(c.Do("LLEN", p.name))

	if err != nil && err.Error() == "redigo: nil returned" {
		//expire
		return 0, nil
	}
	return i, err
}

//enqueue, block if buffer is full
func (p *SimpleRedisQueue) EnqBlocking(data interface{}) {
	p.buffer.EnqBlocking(data)
}

//enqueue, return error if buffer is full
func (p *SimpleRedisQueue) Enq(data interface{}) error {
	return p.buffer.Enq(data)
}

func (p *SimpleRedisQueue) Deq() (interface{}, error) {
	c := p.pool.Get()
	defer c.Close()
	return c.Do("LPOP", p.name)
}

func (p *SimpleRedisQueue) BufferLen() int {
	return p.buffer.Len()
}

func (p *SimpleRedisQueue) BufferCap() int {
	return p.buffer.Cap()
}

func (p *SimpleRedisQueue) PollSize() int {
	return p.pool.ActiveCount()
}
<|endoftext|>


<commit_before>package server

import (
	"github.com/itsankoff/gotcha/common"
)

type FileStore struct {
}

func (f FileStore) AddFile(msg *common.Message) string {
	return ""
}

func (f FileStore) RemoveFile(token string) bool {
	return false
}
<commit_msg>Implement os file store<commit_after>package server

import (
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"
)

type FileStore struct {
	rootFolder  string
	host        string
	networkPath string
}

func NewFileStore(folder string, host string) *FileStore {
	fileStore := &FileStore{
		rootFolder:  folder,
		host:        host,
		networkPath: "/",
	}

	fs := http.FileServer(http.Dir(folder))
	http.Handle(fileStore.networkPath, fs)
	return fileStore
}

func (store FileStore) buildUri(token string) string {
	return store.host + store.networkPath + token
}

func (store FileStore) AddFile(fileContent string) string {
	now := time.Now().UnixNano()
	token := "tmp" + strconv.FormatInt(now, 10)
	filePath := store.rootFolder + "/" + token
	file, err := os.Create(filePath)
	if err != nil {
		log.Println("Failed to open and save", filePath, err)
		return ""
	}

	_, err = file.WriteString(string(fileContent))
	if err != nil {
		log.Println("Failed to write file", filePath, err)
		return ""
	}

	uri := store.buildUri(token)
	log.Println("File saved", uri)

	return uri
}

func (store FileStore) RemoveFile(uri string) bool {
	split := strings.Split(uri, "/")
	token := split[len(split)-1]
	filePath := store.rootFolder + "/" + token
	err := os.Remove(filePath)
	if err != nil {
		log.Println("Failed to remove file for uri", uri)
		return false
	}

	log.Printf("File with %s deleted", uri)
	return true
}
<|endoftext|>


<commit_before>package automaton

import (
	"fmt"
	"github.com/balzaczyy/golucene/core/util"
	"sort"
	"unicode"
)

// util/automaton/Automaton.java

/*
Represents an automaton and all its states and transitions. States
are integers and must be created using {@link #createState}.  Mark a
state as an accept state using {@link #setAccept}.  Add transitions
using {@link #addTransition}.  Each state must have all of its
transitions added at once; if this is too restrictive then use
{@link Automaton.Builder} instead.  State 0 is always the
initial state.  Once a state is finished, either
because you've starting adding transitions to another state or you
call {@link #finishState}, then that states transitions are sorted
(first by min, then max, then dest) and reduced (transitions with
adjacent labels going to the same dest are combined).
*/
type Automaton struct {
	curState      int
	states        []int // 2x
	transitions   []int // 3x
	isAccept      *util.OpenBitSet
	deterministic bool
}

func newEmptyAutomaton() *Automaton {
	return &Automaton{
		deterministic: true,
		curState:      -1,
		isAccept:      util.NewOpenBitSet(),
	}
}

/* Create a new state. */
func (a *Automaton) createState() int {
	state := len(a.states) / 2
	a.states = append(a.states, -1, 0)
	return state
}

/* Set or clear this state as an accept state. */
func (a *Automaton) setAccept(state int, accept bool) {
	assert2(state < a.numStates(), "state=%v is out of bounds (numStates=%v)", state, a.numStates())
	if accept {
		a.isAccept.Set(int64(state))
	} else {
		a.isAccept.Clear(int64(state))
	}
}

/* Returns true if this state is an accept state. */
func (a *Automaton) IsAccept(state int) bool {
	return a.isAccept.Get(int64(state))
}

/* Add a new transition with min = max = label. */
func (a *Automaton) addTransition(source, dest, label int) {
	a.addTransitionRange(source, dest, label, label)
}

/* Add a new transition with the specified source, dest, min, max. */
func (a *Automaton) addTransitionRange(source, dest, min, max int) {
	assert(len(a.transitions)%3 == 0)
	assert2(source < a.numStates(), "source=%v is out of bounds (maxState is %v)", source, a.numStates()-1)
	assert2(dest < a.numStates(), "dest=%v is out of bounds (maxState is %v)", dest, a.numStates()-1)

	if a.curState != source {
		if a.curState != -1 {
			a.finishCurrentState()
		}

		// move to next source:
		a.curState = source
		assert2(a.states[2*a.curState] == -1, "from state (%v) already had transitions added", source)
		assert(a.states[2*a.curState+1] == 0)
		a.states[2*a.curState] = len(a.transitions)
	}

	a.transitions = append(a.transitions, dest, min, max)

	// increment transition count for this state
	a.states[2*a.curState+1]++
}

/*
Add a [virtual] epsilon transition between source and dest. Dest
state must already have all transitions added because this method
simply copies those same transitions over to source.
*/
func (a *Automaton) addEpsilon(source, dest int) {
	t := newTransition()
	count := a.initTransition(dest, t)
	for i := 0; i < count; i++ {
		a.nextTransition(t)
		a.addTransitionRange(source, t.dest, t.min, t.max)
	}
	if a.IsAccept(dest) {
		a.setAccept(source, true)
	}
}

/*
Copies over all state/transition from other. The state numbers are
sequentially assigned (appended).
*/
func (a *Automaton) copy(other *Automaton) {
	// bulk copy and then fixup the state pointers
	stateOffset := a.numStates()
	a.states = append(a.states, other.states...)
	for i := 0; i < len(other.states); i += 2 {
		if a.states[stateOffset*2+i] != -1 {
			a.states[stateOffset*2+i] += len(a.transitions)
		}
	}
	otherNumStates := other.numStates()
	for state := 0; state < otherNumStates; state++ {
		if !other.IsAccept(state) {
			continue
		}
		a.setAccept(stateOffset+state, true)
	}

	// bulk copy and then fixup dest for each transition
	transOffset := len(a.transitions)
	a.transitions = append(a.transitions, other.transitions...)
	for i := 0; i < len(other.transitions); i += 3 {
		a.transitions[transOffset+i] += stateOffset
	}

	if !other.deterministic {
		a.deterministic = false
	}
}

/* Freezes the last state, sorting and reducing the transitions. */
func (a *Automaton) finishCurrentState() {
	numTransitions := a.states[2*a.curState+1]
	assert(numTransitions > 0)

	offset := a.states[2*a.curState]
	start := offset / 3
	util.NewInPlaceMergeSorter(destMinMaxTransitions(a.transitions)).Sort(start, start+numTransitions)

	// reduce any "adjacent" transitions:
	upto, min, max, dest := 0, -1, -1, -1

	for i := 0; i < numTransitions; i++ {
		tDest := a.transitions[offset+3*i]
		tMin := a.transitions[offset+3*i+1]
		tMax := a.transitions[offset+3*i+2]

		if dest == tDest {
			panic("niy")
		} else {
			if dest != -1 {
				a.transitions[offset+3*upto] = dest
				a.transitions[offset+3*upto+1] = min
				a.transitions[offset+3*upto+2] = max
			}
			dest, min, max = tDest, tMin, tMax
		}
	}

	if dest != -1 {
		// last transition
		a.transitions[offset+3*upto] = dest
		a.transitions[offset+3*upto+1] = min
		a.transitions[offset+3*upto+2] = max
		upto++
	}

	a.transitions = a.transitions[:len(a.transitions)-(numTransitions-upto)*3]
	a.states[2*a.curState+1] = upto

	// sort transitions by min/max/dest:
	util.NewInPlaceMergeSorter(minMaxDestTransitions(a.transitions)).Sort(start, start+upto)

	if a.deterministic && upto > 1 {
		panic("niy")
	}
}

/*
Finishes the current state; call this once you are done adding
transitions for a state. This is automatically called if you start
adding transitions to a new source state, but for the last state you
add, you need to call this method yourself.
*/
func (a *Automaton) finishState() {
	if a.curState != -1 {
		a.finishCurrentState()
		a.curState = -1
	}
}

/* How many states this automaton has. */
func (a *Automaton) numStates() int {
	return len(a.states) / 2
}

/* How many transitions this state has. */
func (a *Automaton) numTransitions(state int) int {
	if count := a.states[2*state+1]; count != -1 {
		return count
	}
	return 0
}

type destMinMaxTransitions []int

func (s destMinMaxTransitions) Len() int {
	panic("niy")
}

func (s destMinMaxTransitions) Swap(i, j int) {
	panic("niy")
}

func (s destMinMaxTransitions) Less(i, j int) bool {
	iStart := 3 * i
	jStart := 3 * j

	// first dest:
	iDest := s[iStart]
	jDest := s[jStart]
	if iDest < jDest {
		return true
	} else if iDest > jDest {
		return false
	}

	// then min:
	iMin := s[iStart+1]
	jMin := s[jStart+1]
	if iMin < jMin {
		return true
	} else if iMin > jMin {
		return false
	}

	// then max:
	iMax := s[iStart+2]
	jMax := s[jStart+2]
	return iMax < jMax
}

type minMaxDestTransitions []int

func (s minMaxDestTransitions) Len() int {
	panic("niy")
}

func (s minMaxDestTransitions) Swap(i, j int) {
	panic("niy")
}

func (s minMaxDestTransitions) Less(i, j int) bool {
	panic("niy")
}

/*
Initialize the provided Transition to iterate through all transitions
leaving the specified state. You must call nextTransition() to get
each transition. Returns the number of transitions leaving this tate.
*/
func (a *Automaton) initTransition(state int, t *Transition) int {
	assert2(state < a.numStates(), "state=%v nextState=%v", state, a.numStates())
	t.source = state
	t.transitionUpto = a.states[2*state]
	return a.numTransitions(state)
}

/* Iterate to the next transition after the provided one */
func (a *Automaton) nextTransition(t *Transition) {
	// make sure there is still a transition left
	assert((t.transitionUpto + 3 - a.states[2*t.source]) <= 3*a.states[2*t.source+1])
	t.dest = a.transitions[t.transitionUpto]
	t.min = a.transitions[t.transitionUpto+1]
	t.max = a.transitions[t.transitionUpto+2]
	t.transitionUpto += 3
}

/*
Fill the provided Transition with the index'th transition leaving the
specified state.
*/
func (a *Automaton) transition(state, index int, t *Transition) {
	panic("niy")
}

// L563
/* Returns sorted array of all interval start points. */
func (a *Automaton) startPoints() []int {
	pointset := make(map[int]bool)
	pointset[MIN_CODE_POINT] = true
	fmt.Println("getStartPoints")
	for s := 0; s < len(a.states); s += 2 {
		trans := a.states[s]
		limit := trans + 3*a.states[s+1]
		fmt.Printf("  state=%v trans=%v limit=%v\n", s/2, trans, limit)
		for trans < limit {
			min, max := a.transitions[trans+1], a.transitions[trans+2]
			fmt.Printf("    min=%v\n", min)
			pointset[min] = true
			if max < unicode.MaxRune {
				pointset[max+1] = true
			}
			trans += 3
		}
	}
	var points []int
	for m, _ := range pointset {
		points = append(points, m)
	}
	sort.Ints(points)
	return points
}

/* Performs lookup in transitions, assuming determinism. */
func (a *Automaton) step(state, label int) int {
	panic("niy")
}

// Go doesn't have unicode.MinRune which should be 0
const MIN_CODE_POINT = 0

type AutomatonBuilder struct {
	transitions []int
	a           *Automaton
}

func newAutomatonBuilder() *AutomatonBuilder {
	return &AutomatonBuilder{
		a: newEmptyAutomaton(),
	}
}

func (b *AutomatonBuilder) addTransitionRange(source, dest, min, max int) {
	b.transitions = append(b.transitions, source, dest, min, max)
}

type srcMinMaxDestSorter []int

func (s srcMinMaxDestSorter) Len() int {
	panic("niy")
}

func (s srcMinMaxDestSorter) Swap(i, j int) {
	iStart, jStart := 4*i, 4*j
	for n := 0; n < 4; n++ {
		s[iStart+n], s[jStart+n] = s[jStart+n], s[iStart+n]
	}
}

func (s srcMinMaxDestSorter) Less(i, j int) bool {
	iStart, jStart := 4*i, 4*j

	iSrc, jSrc := s[iStart], s[jStart]
	if iSrc < jSrc {
		return true
	} else if iSrc > jSrc {
		return false
	}

	iMin, jMin := s[iStart+2], s[jStart+2]
	if iMin < jMin {
		return true
	} else if iMin > jMin {
		return false
	}

	iMax, jMax := s[iStart+3], s[jStart+3]
	if iMax < jMax {
		return true
	} else if iMax > jMax {
		return false
	}

	iDest, jDest := s[iStart+1], s[jStart+1]
	return iDest < jDest
}

/* Compiles all added states and transitions into a new Automaton and returns it. */
func (b *AutomatonBuilder) finish() *Automaton {
	fmt.Printf("LA.Builder.finish: count=%v\n", len(b.transitions)/4)
	fmt.Println("finish pending")
	util.NewInPlaceMergeSorter(srcMinMaxDestSorter(b.transitions)).Sort(0, len(b.transitions)/4)
	for upto := 0; upto < len(b.transitions); upto += 4 {
		b.a.addTransitionRange(
			b.transitions[upto],
			b.transitions[upto+1],
			b.transitions[upto+2],
			b.transitions[upto]+3,
		)
	}

	b.a.finishState()
	return b.a
}

func (b *AutomatonBuilder) createState() int {
	return b.a.createState()
}

func (b *AutomatonBuilder) setAccept(state int, accept bool) {
	b.a.setAccept(state, accept)
}

func (b *AutomatonBuilder) isAccept(state int) bool {
	return b.a.IsAccept(state)
}
<commit_msg>implement Automaton.step()<commit_after>package automaton

import (
	"fmt"
	"github.com/balzaczyy/golucene/core/util"
	"sort"
	"unicode"
)

// util/automaton/Automaton.java

/*
Represents an automaton and all its states and transitions. States
are integers and must be created using {@link #createState}.  Mark a
state as an accept state using {@link #setAccept}.  Add transitions
using {@link #addTransition}.  Each state must have all of its
transitions added at once; if this is too restrictive then use
{@link Automaton.Builder} instead.  State 0 is always the
initial state.  Once a state is finished, either
because you've starting adding transitions to another state or you
call {@link #finishState}, then that states transitions are sorted
(first by min, then max, then dest) and reduced (transitions with
adjacent labels going to the same dest are combined).
*/
type Automaton struct {
	curState      int
	states        []int // 2x
	transitions   []int // 3x
	isAccept      *util.OpenBitSet
	deterministic bool
}

func newEmptyAutomaton() *Automaton {
	return &Automaton{
		deterministic: true,
		curState:      -1,
		isAccept:      util.NewOpenBitSet(),
	}
}

/* Create a new state. */
func (a *Automaton) createState() int {
	state := len(a.states) / 2
	a.states = append(a.states, -1, 0)
	return state
}

/* Set or clear this state as an accept state. */
func (a *Automaton) setAccept(state int, accept bool) {
	assert2(state < a.numStates(), "state=%v is out of bounds (numStates=%v)", state, a.numStates())
	if accept {
		a.isAccept.Set(int64(state))
	} else {
		a.isAccept.Clear(int64(state))
	}
}

/* Returns true if this state is an accept state. */
func (a *Automaton) IsAccept(state int) bool {
	return a.isAccept.Get(int64(state))
}

/* Add a new transition with min = max = label. */
func (a *Automaton) addTransition(source, dest, label int) {
	a.addTransitionRange(source, dest, label, label)
}

/* Add a new transition with the specified source, dest, min, max. */
func (a *Automaton) addTransitionRange(source, dest, min, max int) {
	assert(len(a.transitions)%3 == 0)
	assert2(source < a.numStates(), "source=%v is out of bounds (maxState is %v)", source, a.numStates()-1)
	assert2(dest < a.numStates(), "dest=%v is out of bounds (maxState is %v)", dest, a.numStates()-1)

	if a.curState != source {
		if a.curState != -1 {
			a.finishCurrentState()
		}

		// move to next source:
		a.curState = source
		assert2(a.states[2*a.curState] == -1, "from state (%v) already had transitions added", source)
		assert(a.states[2*a.curState+1] == 0)
		a.states[2*a.curState] = len(a.transitions)
	}

	a.transitions = append(a.transitions, dest, min, max)

	// increment transition count for this state
	a.states[2*a.curState+1]++
}

/*
Add a [virtual] epsilon transition between source and dest. Dest
state must already have all transitions added because this method
simply copies those same transitions over to source.
*/
func (a *Automaton) addEpsilon(source, dest int) {
	t := newTransition()
	count := a.initTransition(dest, t)
	for i := 0; i < count; i++ {
		a.nextTransition(t)
		a.addTransitionRange(source, t.dest, t.min, t.max)
	}
	if a.IsAccept(dest) {
		a.setAccept(source, true)
	}
}

/*
Copies over all state/transition from other. The state numbers are
sequentially assigned (appended).
*/
func (a *Automaton) copy(other *Automaton) {
	// bulk copy and then fixup the state pointers
	stateOffset := a.numStates()
	a.states = append(a.states, other.states...)
	for i := 0; i < len(other.states); i += 2 {
		if a.states[stateOffset*2+i] != -1 {
			a.states[stateOffset*2+i] += len(a.transitions)
		}
	}
	otherNumStates := other.numStates()
	for state := 0; state < otherNumStates; state++ {
		if !other.IsAccept(state) {
			continue
		}
		a.setAccept(stateOffset+state, true)
	}

	// bulk copy and then fixup dest for each transition
	transOffset := len(a.transitions)
	a.transitions = append(a.transitions, other.transitions...)
	for i := 0; i < len(other.transitions); i += 3 {
		a.transitions[transOffset+i] += stateOffset
	}

	if !other.deterministic {
		a.deterministic = false
	}
}

/* Freezes the last state, sorting and reducing the transitions. */
func (a *Automaton) finishCurrentState() {
	numTransitions := a.states[2*a.curState+1]
	assert(numTransitions > 0)

	offset := a.states[2*a.curState]
	start := offset / 3
	util.NewInPlaceMergeSorter(destMinMaxTransitions(a.transitions)).Sort(start, start+numTransitions)

	// reduce any "adjacent" transitions:
	upto, min, max, dest := 0, -1, -1, -1

	for i := 0; i < numTransitions; i++ {
		tDest := a.transitions[offset+3*i]
		tMin := a.transitions[offset+3*i+1]
		tMax := a.transitions[offset+3*i+2]

		if dest == tDest {
			panic("niy")
		} else {
			if dest != -1 {
				a.transitions[offset+3*upto] = dest
				a.transitions[offset+3*upto+1] = min
				a.transitions[offset+3*upto+2] = max
			}
			dest, min, max = tDest, tMin, tMax
		}
	}

	if dest != -1 {
		// last transition
		a.transitions[offset+3*upto] = dest
		a.transitions[offset+3*upto+1] = min
		a.transitions[offset+3*upto+2] = max
		upto++
	}

	a.transitions = a.transitions[:len(a.transitions)-(numTransitions-upto)*3]
	a.states[2*a.curState+1] = upto

	// sort transitions by min/max/dest:
	util.NewInPlaceMergeSorter(minMaxDestTransitions(a.transitions)).Sort(start, start+upto)

	if a.deterministic && upto > 1 {
		panic("niy")
	}
}

/*
Finishes the current state; call this once you are done adding
transitions for a state. This is automatically called if you start
adding transitions to a new source state, but for the last state you
add, you need to call this method yourself.
*/
func (a *Automaton) finishState() {
	if a.curState != -1 {
		a.finishCurrentState()
		a.curState = -1
	}
}

/* How many states this automaton has. */
func (a *Automaton) numStates() int {
	return len(a.states) / 2
}

/* How many transitions this state has. */
func (a *Automaton) numTransitions(state int) int {
	if count := a.states[2*state+1]; count != -1 {
		return count
	}
	return 0
}

type destMinMaxTransitions []int

func (s destMinMaxTransitions) Len() int {
	panic("niy")
}

func (s destMinMaxTransitions) Swap(i, j int) {
	panic("niy")
}

func (s destMinMaxTransitions) Less(i, j int) bool {
	iStart := 3 * i
	jStart := 3 * j

	// first dest:
	iDest := s[iStart]
	jDest := s[jStart]
	if iDest < jDest {
		return true
	} else if iDest > jDest {
		return false
	}

	// then min:
	iMin := s[iStart+1]
	jMin := s[jStart+1]
	if iMin < jMin {
		return true
	} else if iMin > jMin {
		return false
	}

	// then max:
	iMax := s[iStart+2]
	jMax := s[jStart+2]
	return iMax < jMax
}

type minMaxDestTransitions []int

func (s minMaxDestTransitions) Len() int {
	panic("niy")
}

func (s minMaxDestTransitions) Swap(i, j int) {
	panic("niy")
}

func (s minMaxDestTransitions) Less(i, j int) bool {
	panic("niy")
}

/*
Initialize the provided Transition to iterate through all transitions
leaving the specified state. You must call nextTransition() to get
each transition. Returns the number of transitions leaving this tate.
*/
func (a *Automaton) initTransition(state int, t *Transition) int {
	assert2(state < a.numStates(), "state=%v nextState=%v", state, a.numStates())
	t.source = state
	t.transitionUpto = a.states[2*state]
	return a.numTransitions(state)
}

/* Iterate to the next transition after the provided one */
func (a *Automaton) nextTransition(t *Transition) {
	// make sure there is still a transition left
	assert((t.transitionUpto + 3 - a.states[2*t.source]) <= 3*a.states[2*t.source+1])
	t.dest = a.transitions[t.transitionUpto]
	t.min = a.transitions[t.transitionUpto+1]
	t.max = a.transitions[t.transitionUpto+2]
	t.transitionUpto += 3
}

/*
Fill the provided Transition with the index'th transition leaving the
specified state.
*/
func (a *Automaton) transition(state, index int, t *Transition) {
	panic("niy")
}

// L563
/* Returns sorted array of all interval start points. */
func (a *Automaton) startPoints() []int {
	pointset := make(map[int]bool)
	pointset[MIN_CODE_POINT] = true
	fmt.Println("getStartPoints")
	for s := 0; s < len(a.states); s += 2 {
		trans := a.states[s]
		limit := trans + 3*a.states[s+1]
		fmt.Printf("  state=%v trans=%v limit=%v\n", s/2, trans, limit)
		for trans < limit {
			min, max := a.transitions[trans+1], a.transitions[trans+2]
			fmt.Printf("    min=%v\n", min)
			pointset[min] = true
			if max < unicode.MaxRune {
				pointset[max+1] = true
			}
			trans += 3
		}
	}
	var points []int
	for m, _ := range pointset {
		points = append(points, m)
	}
	sort.Ints(points)
	return points
}

/* Performs lookup in transitions, assuming determinism. */
func (a *Automaton) step(state, label int) int {
	assert(state >= 0)
	assert(label >= 0)
	trans := a.states[2*state]
	limit := trans + 3*a.states[2*state+1]
	// TODO binary search
	for trans < limit {
		dest, min, max := a.transitions[trans], a.transitions[trans+1], a.transitions[trans+2]
		if min <= label && label <= max {
			return dest
		}
		trans += 3
	}
	return -1
}

// Go doesn't have unicode.MinRune which should be 0
const MIN_CODE_POINT = 0

type AutomatonBuilder struct {
	transitions []int
	a           *Automaton
}

func newAutomatonBuilder() *AutomatonBuilder {
	return &AutomatonBuilder{
		a: newEmptyAutomaton(),
	}
}

func (b *AutomatonBuilder) addTransitionRange(source, dest, min, max int) {
	b.transitions = append(b.transitions, source, dest, min, max)
}

type srcMinMaxDestSorter []int

func (s srcMinMaxDestSorter) Len() int {
	panic("niy")
}

func (s srcMinMaxDestSorter) Swap(i, j int) {
	iStart, jStart := 4*i, 4*j
	for n := 0; n < 4; n++ {
		s[iStart+n], s[jStart+n] = s[jStart+n], s[iStart+n]
	}
}

func (s srcMinMaxDestSorter) Less(i, j int) bool {
	iStart, jStart := 4*i, 4*j

	iSrc, jSrc := s[iStart], s[jStart]
	if iSrc < jSrc {
		return true
	} else if iSrc > jSrc {
		return false
	}

	iMin, jMin := s[iStart+2], s[jStart+2]
	if iMin < jMin {
		return true
	} else if iMin > jMin {
		return false
	}

	iMax, jMax := s[iStart+3], s[jStart+3]
	if iMax < jMax {
		return true
	} else if iMax > jMax {
		return false
	}

	iDest, jDest := s[iStart+1], s[jStart+1]
	return iDest < jDest
}

/* Compiles all added states and transitions into a new Automaton and returns it. */
func (b *AutomatonBuilder) finish() *Automaton {
	fmt.Printf("LA.Builder.finish: count=%v\n", len(b.transitions)/4)
	fmt.Println("finish pending")
	util.NewInPlaceMergeSorter(srcMinMaxDestSorter(b.transitions)).Sort(0, len(b.transitions)/4)
	for upto := 0; upto < len(b.transitions); upto += 4 {
		b.a.addTransitionRange(
			b.transitions[upto],
			b.transitions[upto+1],
			b.transitions[upto+2],
			b.transitions[upto]+3,
		)
	}

	b.a.finishState()
	return b.a
}

func (b *AutomatonBuilder) createState() int {
	return b.a.createState()
}

func (b *AutomatonBuilder) setAccept(state int, accept bool) {
	b.a.setAccept(state, accept)
}

func (b *AutomatonBuilder) isAccept(state int) bool {
	return b.a.IsAccept(state)
}
<|endoftext|>


<commit_before>package service

import (
	"encoding/json"
	"fmt"
	"time"

	"github.com/kyleterry/tenyks/irc"
	. "github.com/kyleterry/tenyks/version"
)

type servicefn func(*Connection, *Message)

func (self *ServiceEngine) AddHandler(name string, fn servicefn) {
	handler := irc.NewHandler(func(p ...interface{}) {
		fn(p[0].(*Connection), p[1].(*Message))
	})
	self.CommandRg.AddHandler(name, handler)
}

func (self *ServiceEngine) addBaseHandlers() {
	self.AddHandler("PRIVMSG", (*Connection).PrivmsgServiceHandler)
	self.AddHandler("REGISTER", (*Connection).RegisterServiceHandler)
	self.AddHandler("BYE", (*Connection).ByeServiceHandler)
	self.AddHandler("PONG", (*Connection).PongServiceHandler)
}

func (self *Connection) PrivmsgIrcHandler(conn *irc.Connection, msg *irc.Message) {
	serviceMsg := Message{}
	serviceMsg.Target = msg.Params[0]
	serviceMsg.Command = msg.Command
	serviceMsg.Mask = msg.Host
	serviceMsg.Direct = irc.IsDirect(msg.Trail, conn.GetCurrentNick())
	serviceMsg.Nick = msg.Nick
	serviceMsg.Host = msg.Host
	serviceMsg.Full_message = msg.RawMsg
	serviceMsg.User = msg.Ident
	serviceMsg.From_channel = irc.IsChannel(msg.Params[0])
	serviceMsg.Connection = conn.Name
	serviceMsg.Meta = &Meta{"Tenyks", TenyksVersion, nil}
	if serviceMsg.Direct {
		serviceMsg.Payload = irc.StripNickOnDirect(msg.Trail, conn.GetCurrentNick())
	} else {
		serviceMsg.Payload = msg.Trail
	}

	jsonBytes, err := json.Marshal(serviceMsg)
	if err != nil {
		log.Fatal(err)
	}
	self.Out <- string(jsonBytes[:])
}

func (self *Connection) PrivmsgServiceHandler(msg *Message) {
	conn := self.getIrcConnByName(msg.Connection)
	if conn != nil {
		msgStr := fmt.Sprintf("%s %s :%s", msg.Command, msg.Target, msg.Payload)
		conn.Out <- msgStr
	} else {
		log.Debug("[service] No such connection `%s`. Ignoring.",
			msg.Connection)
	}
}

func (self *Connection) RegisterServiceHandler(msg *Message) {
	meta := msg.Meta
	if meta.SID == nil {
		log.Error("[service] ERROR: UUID required to register with Tenyks")
		return
	}
	log.Debug("[service] %s (%s) wants to register", meta.SID.UUID.String(), meta.Name)
	srv := &Service{}
	srv.Name = meta.Name
	srv.Version = meta.Version
	srv.Online = true
	srv.LastPing = time.Now()
	srv.UUID = meta.SID.UUID
	self.engine.ServiceRg.RegisterService(srv)
}

func (self *Connection) ByeServiceHandler(msg *Message) {
	meta := msg.Meta
	log.Debug("[service] %s (%s) is hanging up", meta.SID.UUID.String(), meta.Name)
	srv := self.engine.ServiceRg.GetServiceByUUID(meta.Name)
	if srv != nil {
		srv.Online = false
	}
}

type ServiceListMessage struct {
	Services map[string]*Service `json:"services"`
	Command  string              `json:"command"`
	Meta     *Meta               `json:"meta"`
}

func (self *Connection) ListServiceHandler(msg *Message) {
	serviceList := &ServiceListMessage{}
	serviceList.Services = self.engine.ServiceRg.services
	serviceList.Command = "SERVICES"
	serviceList.Meta = &Meta{"Tenyks", TenyksVersion, nil}
	jsonBytes, err := json.Marshal(serviceList)
	if err != nil {
		log.Fatal(err)
	}
	self.Out <- string(jsonBytes[:])
}

const (
	ServiceOnline = true
	ServiceOffline = false
)

func (self *Connection) PingServices() {
	log.Debug("[service] Starting pinger")
	for {
		<-time.After(time.Second * 120)
		log.Debug("[service] PINGing services")
		msg := &Message{
			Command: "PING",
			Payload: "!tenyks",
		}
		jsonBytes, err := json.Marshal(msg)
		if err != nil {
			log.Error("Cannot marshal PING message")
			continue
		}
		self.Out <- string(jsonBytes[:])

		services := self.engine.ServiceRg.services
		for _, service := range services {
			if service.Online {
				service.LastPing = time.Now()
			}
		}
	}
}

func (self *Connection) PongServiceHandler(msg *Message) {
	meta := msg.Meta
	self.engine.UpdateService(meta.SID.UUID.String(), ServiceOnline)
}
<commit_msg>Gracefully handle services without UUIDs<commit_after>package service

import (
	"encoding/json"
	"fmt"
	"time"

	"github.com/kyleterry/tenyks/irc"
	. "github.com/kyleterry/tenyks/version"
)

type servicefn func(*Connection, *Message)

func (self *ServiceEngine) AddHandler(name string, fn servicefn) {
	handler := irc.NewHandler(func(p ...interface{}) {
		fn(p[0].(*Connection), p[1].(*Message))
	})
	self.CommandRg.AddHandler(name, handler)
}

func (self *ServiceEngine) addBaseHandlers() {
	self.AddHandler("PRIVMSG", (*Connection).PrivmsgServiceHandler)
	self.AddHandler("REGISTER", (*Connection).RegisterServiceHandler)
	self.AddHandler("BYE", (*Connection).ByeServiceHandler)
	self.AddHandler("PONG", (*Connection).PongServiceHandler)
}

func (self *Connection) PrivmsgIrcHandler(conn *irc.Connection, msg *irc.Message) {
	serviceMsg := Message{}
	serviceMsg.Target = msg.Params[0]
	serviceMsg.Command = msg.Command
	serviceMsg.Mask = msg.Host
	serviceMsg.Direct = irc.IsDirect(msg.Trail, conn.GetCurrentNick())
	serviceMsg.Nick = msg.Nick
	serviceMsg.Host = msg.Host
	serviceMsg.Full_message = msg.RawMsg
	serviceMsg.User = msg.Ident
	serviceMsg.From_channel = irc.IsChannel(msg.Params[0])
	serviceMsg.Connection = conn.Name
	serviceMsg.Meta = &Meta{"Tenyks", TenyksVersion, nil}
	if serviceMsg.Direct {
		serviceMsg.Payload = irc.StripNickOnDirect(msg.Trail, conn.GetCurrentNick())
	} else {
		serviceMsg.Payload = msg.Trail
	}

	jsonBytes, err := json.Marshal(serviceMsg)
	if err != nil {
		log.Fatal(err)
	}
	self.Out <- string(jsonBytes[:])
}

func (self *Connection) PrivmsgServiceHandler(msg *Message) {
	conn := self.getIrcConnByName(msg.Connection)
	if conn != nil {
		msgStr := fmt.Sprintf("%s %s :%s", msg.Command, msg.Target, msg.Payload)
		conn.Out <- msgStr
	} else {
		log.Debug("[service] No such connection `%s`. Ignoring.",
			msg.Connection)
	}
}

func (self *Connection) RegisterServiceHandler(msg *Message) {
	meta := msg.Meta
	if meta.SID == nil || meta.SID.UUID == nil {
		log.Error("[service] ERROR: UUID required to register with Tenyks")
		return
	}
	log.Debug("[service] %s (%s) wants to register", meta.SID.UUID.String(), meta.Name)
	srv := &Service{}
	srv.Name = meta.Name
	srv.Version = meta.Version
	srv.Online = true
	srv.LastPing = time.Now()
	srv.UUID = meta.SID.UUID
	self.engine.ServiceRg.RegisterService(srv)
}

func (self *Connection) ByeServiceHandler(msg *Message) {
	meta := msg.Meta
	if meta.SID.UUID != nil {
		log.Debug("[service] %s (%s) is hanging up", meta.SID.UUID.String(), meta.Name)
		srv := self.engine.ServiceRg.GetServiceByUUID(meta.Name)
		if srv != nil {
			srv.Online = false
		}
	}
}

type ServiceListMessage struct {
	Services map[string]*Service `json:"services"`
	Command  string              `json:"command"`
	Meta     *Meta               `json:"meta"`
}

func (self *Connection) ListServiceHandler(msg *Message) {
	serviceList := &ServiceListMessage{}
	serviceList.Services = self.engine.ServiceRg.services
	serviceList.Command = "SERVICES"
	serviceList.Meta = &Meta{"Tenyks", TenyksVersion, nil}
	jsonBytes, err := json.Marshal(serviceList)
	if err != nil {
		log.Fatal(err)
	}
	self.Out <- string(jsonBytes[:])
}

const (
	ServiceOnline = true
	ServiceOffline = false
)

func (self *Connection) PingServices() {
	log.Debug("[service] Starting pinger")
	for {
		<-time.After(time.Second * 120)
		log.Debug("[service] PINGing services")
		msg := &Message{
			Command: "PING",
			Payload: "!tenyks",
		}
		jsonBytes, err := json.Marshal(msg)
		if err != nil {
			log.Error("Cannot marshal PING message")
			continue
		}
		self.Out <- string(jsonBytes[:])

		services := self.engine.ServiceRg.services
		for _, service := range services {
			if service.Online {
				service.LastPing = time.Now()
			}
		}
	}
}

func (self *Connection) PongServiceHandler(msg *Message) {
	meta := msg.Meta
	if meta.SID != nil && meta.SID.UUID != nil {
		self.engine.UpdateService(meta.SID.UUID.String(), ServiceOnline)
	}
}
<|endoftext|>


<commit_before>package vfs

import (
	"fmt"
	"github.com/dotcloud/docker/daemon/graphdriver"
	"os"
	"os/exec"
	"path"
)

func init() {
	graphdriver.Register("vfs", Init)
}

func Init(home string, options []string) (graphdriver.Driver, error) {
	d := &Driver{
		home: home,
	}
	return d, nil
}

type Driver struct {
	home string
}

func (d *Driver) String() string {
	return "vfs"
}

func (d *Driver) Status() [][2]string {
	return nil
}

func (d *Driver) Cleanup() error {
	return nil
}

func copyDir(src, dst string) error {
	if output, err := exec.Command("cp", "-aT", "--reflink=auto", src, dst).CombinedOutput(); err != nil {
		return fmt.Errorf("Error VFS copying directory: %s (%s)", err, output)
	}
	return nil
}

func (d *Driver) Create(id, parent string) error {
	dir := d.dir(id)
	if err := os.MkdirAll(path.Dir(dir), 0700); err != nil {
		return err
	}
	if err := os.Mkdir(dir, 0755); err != nil {
		return err
	}
	if parent == "" {
		return nil
	}
	parentDir, err := d.Get(parent, "")
	if err != nil {
		return fmt.Errorf("%s: %s", parent, err)
	}
	if err := copyDir(parentDir, dir); err != nil {
		return err
	}
	return nil
}

func (d *Driver) dir(id string) string {
	return path.Join(d.home, "dir", path.Base(id))
}

func (d *Driver) Remove(id string) error {
	if _, err := os.Stat(d.dir(id)); err != nil {
		return err
	}
	return os.RemoveAll(d.dir(id))
}

func (d *Driver) Get(id, mountLabel string) (string, error) {
	dir := d.dir(id)
	if st, err := os.Stat(dir); err != nil {
		return "", err
	} else if !st.IsDir() {
		return "", fmt.Errorf("%s: not a directory", dir)
	}
	return dir, nil
}

func (d *Driver) Put(id string) {
	// The vfs driver has no runtime resources (e.g. mounts)
	// to clean up, so we don't need anything here
}

func (d *Driver) Exists(id string) bool {
	_, err := os.Stat(d.dir(id))
	return err == nil
}
<commit_msg>cp's -T and --reflink=auto are only available on GNU coreutils<commit_after>package vfs

import (
	"bytes"
	"fmt"
	"github.com/dotcloud/docker/daemon/graphdriver"
	"os"
	"os/exec"
	"path"
)

func init() {
	graphdriver.Register("vfs", Init)
}

func Init(home string, options []string) (graphdriver.Driver, error) {
	d := &Driver{
		home: home,
	}
	return d, nil
}

type Driver struct {
	home string
}

func (d *Driver) String() string {
	return "vfs"
}

func (d *Driver) Status() [][2]string {
	return nil
}

func (d *Driver) Cleanup() error {
	return nil
}

func isGNUcoreutils() bool {
	if stdout, err := exec.Command("cp", "--version").Output(); err == nil {
		return bytes.Contains(stdout, []byte("GNU coreutils"))
	}

	return false
}

func copyDir(src, dst string) error {
	argv := make([]string, 0, 4)

	if isGNUcoreutils() {
		argv = append(argv, "-aT", "--reflink=auto", src, dst)
	} else {
		argv = append(argv, "-a", src+"/.", dst+"/.")
	}

	if output, err := exec.Command("cp", argv...).CombinedOutput(); err != nil {
		return fmt.Errorf("Error VFS copying directory: %s (%s)", err, output)
	}
	return nil
}

func (d *Driver) Create(id, parent string) error {
	dir := d.dir(id)
	if err := os.MkdirAll(path.Dir(dir), 0700); err != nil {
		return err
	}
	if err := os.Mkdir(dir, 0755); err != nil {
		return err
	}
	if parent == "" {
		return nil
	}
	parentDir, err := d.Get(parent, "")
	if err != nil {
		return fmt.Errorf("%s: %s", parent, err)
	}
	if err := copyDir(parentDir, dir); err != nil {
		return err
	}
	return nil
}

func (d *Driver) dir(id string) string {
	return path.Join(d.home, "dir", path.Base(id))
}

func (d *Driver) Remove(id string) error {
	if _, err := os.Stat(d.dir(id)); err != nil {
		return err
	}
	return os.RemoveAll(d.dir(id))
}

func (d *Driver) Get(id, mountLabel string) (string, error) {
	dir := d.dir(id)
	if st, err := os.Stat(dir); err != nil {
		return "", err
	} else if !st.IsDir() {
		return "", fmt.Errorf("%s: not a directory", dir)
	}
	return dir, nil
}

func (d *Driver) Put(id string) {
	// The vfs driver has no runtime resources (e.g. mounts)
	// to clean up, so we don't need anything here
}

func (d *Driver) Exists(id string) bool {
	_, err := os.Stat(d.dir(id))
	return err == nil
}
<|endoftext|>


<commit_before>/*
 * MumbleDJ
 * By Matthieu Grieger
 * services/youtube.go
 * Copyright (c) 2016 Matthieu Grieger (MIT License)
 */

package services

import (
	"errors"
	"fmt"
	"math"
	"net/http"
	"regexp"

	"github.com/ChannelMeter/iso8601duration"
	"github.com/antonholmquist/jason"
	"github.com/layeh/gumble/gumble"
	"github.com/matthieugrieger/mumbledj/bot"
	"github.com/matthieugrieger/mumbledj/interfaces"
	"github.com/spf13/viper"
)

// YouTube is a wrapper around the YouTube Data API.
// https://developers.google.com/youtube/v3/docs/
type YouTube struct {
	*GenericService
}

// NewYouTubeService returns an initialized YouTube service object.
func NewYouTubeService() *YouTube {
	return &YouTube{
		&GenericService{
			ReadableName: "YouTube",
			Format:       "bestaudio",
			TrackRegex: []*regexp.Regexp{
				regexp.MustCompile(`https?:\/\/www.youtube.com\/watch\?v=(?P<id>[\w-]+)(?P<timestamp>\&t=\d*m?\d*s?)?`),
				regexp.MustCompile(`https?:\/\/youtube.com\/watch\?v=(?P<id>[\w-]+)(?P<timestamp>\&t=\d*m?\d*s?)?`),
				regexp.MustCompile(`https?:\/\/youtu.be\/(?P<id>[\w-]+)(?P<timestamp>\?t=\d*m?\d*s?)?`),
				regexp.MustCompile(`https?:\/\/youtube.com\/v\/(?P<id>[\w-]+)(?P<timestamp>\?t=\d*m?\d*s?)?`),
				regexp.MustCompile(`https?:\/\/www.youtube.com\/v\/(?P<id>[\w-]+)(?P<timestamp>\?t=\d*m?\d*s?)?`),
			},
			PlaylistRegex: []*regexp.Regexp{
				regexp.MustCompile(`https?:\/\/www\.youtube\.com\/playlist\?list=(?P<id>[\w-]+)`),
			},
		},
	}
}

// CheckAPIKey performs a test API call with the API key
// provided in the configuration file to determine if the
// service should be enabled.
func (yt *YouTube) CheckAPIKey() error {
	var (
		response *http.Response
		v        *jason.Object
		err      error
	)

	if viper.GetString("api_keys.youtube") == "" {
		return errors.New("No YouTube API key has been provided")
	}
	url := "https://www.googleapis.com/youtube/v3/videos?part=snippet&id=KQY9zrjPBjo&key=%s"
	response, err = http.Get(fmt.Sprintf(url, viper.GetString("api_keys.youtube")))
	defer response.Body.Close()
	if err != nil {
		return err
	}

	if v, err = jason.NewObjectFromReader(response.Body); err != nil {
		return err
	}

	if v, err = v.GetObject("error"); err == nil {
		message, _ := v.GetString("message")
		code, _ := v.GetInt64("code")
		errArray, _ := v.GetObjectArray("errors")
		reason, _ := errArray[0].GetString("reason")

		return fmt.Errorf("%d: %s (reason: %s)", code, message, reason)
	}
	return nil
}

// GetTracks uses the passed URL to find and return
// tracks associated with the URL. An error is returned
// if any error occurs during the API call.
func (yt *YouTube) GetTracks(url string, submitter *gumble.User) ([]interfaces.Track, error) {
	var (
		playlistURL      string
		playlistItemsURL string
		id               string
		err              error
		resp             *http.Response
		v                *jason.Object
		track            bot.Track
		tracks           []interfaces.Track
	)

	playlistURL = "https://www.googleapis.com/youtube/v3/playlists?part=snippet&id=%s&key=%s"
	playlistItemsURL = "https://www.googleapis.com/youtube/v3/playlistItems?part=snippet,contentDetails&playlistId=%s&maxResults=%d&key=%s&pageToken=%s"
	id, err = yt.getID(url)
	if err != nil {
		return nil, err
	}

	if yt.isPlaylist(url) {
		resp, err = http.Get(fmt.Sprintf(playlistURL, id, viper.GetString("api_keys.youtube")))
		defer resp.Body.Close()
		if err != nil {
			return nil, err
		}

		v, err = jason.NewObjectFromReader(resp.Body)
		if err != nil {
			return nil, err
		}

		items, _ := v.GetObjectArray("items")
		item := items[0]

		title, _ := item.GetString("snippet", "title")

		playlist := &bot.Playlist{
			ID:        id,
			Title:     title,
			Submitter: submitter.Name,
			Service:   yt.ReadableName,
		}

		maxItems := math.MaxInt32
		if viper.GetInt("queue.max_tracks_per_playlist") > 0 {
			maxItems = viper.GetInt("queue.max_tracks_per_playlist")
		}

		// YouTube playlist searches return a max of 50 results per page
		maxResults := 50
		if maxResults > maxItems {
			maxResults = maxItems
		}

		pageToken := ""
		for len(tracks) < maxItems {
			curResp, curErr := http.Get(fmt.Sprintf(playlistItemsURL, id, maxResults, viper.GetString("api_keys.youtube"), pageToken))
			defer curResp.Body.Close()
			if curErr != nil {
				// An error occurred, simply skip this track.
				continue
			}

			v, err = jason.NewObjectFromReader(curResp.Body)
			if err != nil {
				// An error occurred, simply skip this track.
				continue
			}

			curTracks, _ := v.GetObjectArray("items")
			for _, track := range curTracks {
				videoID, _ := track.GetString("snippet", "resourceId", "videoId")

				// Unfortunately we have to execute another API call for each video as the YouTube API does not
				// return video durations from the playlistItems endpoint...
				newTrack, _ := yt.getTrack(videoID, submitter)
				newTrack.Playlist = playlist
				tracks = append(tracks, newTrack)

				if len(tracks) >= maxItems {
					break
				}
			}

			pageToken, _ = v.GetString("nextPageToken")
			if pageToken == "" {
				break
			}
		}

		if len(tracks) == 0 {
			return nil, errors.New("Invalid playlist. No tracks were added")
		}
		return tracks, nil
	}

	track, err = yt.getTrack(id, submitter)
	if err != nil {
		return nil, err
	}
	tracks = append(tracks, track)
	return tracks, nil
}

func (yt *YouTube) getTrack(id string, submitter *gumble.User) (bot.Track, error) {
	var (
		resp *http.Response
		err  error
		v    *jason.Object
	)

	videoURL := "https://www.googleapis.com/youtube/v3/videos?part=snippet,contentDetails&id=%s&key=%s"
	resp, err = http.Get(fmt.Sprintf(videoURL, id, viper.GetString("api_keys.youtube")))
	defer resp.Body.Close()
	if err != nil {
		return bot.Track{}, err
	}

	v, err = jason.NewObjectFromReader(resp.Body)
	if err != nil {
		return bot.Track{}, err
	}
	items, _ := v.GetObjectArray("items")
	item := items[0]
	title, _ := item.GetString("snippet", "title")
	thumbnail, _ := item.GetString("snippet", "thumbnails", "high", "url")
	author, _ := item.GetString("snippet", "channelTitle")
	durationString, _ := item.GetString("contentDetails", "duration")
	durationConverted, _ := duration.FromString(durationString)
	duration := durationConverted.ToDuration()

	return bot.Track{
		ID:           id,
		URL:          "https://youtube.com/watch?v=" + id,
		Title:        title,
		Author:       author,
		Submitter:    submitter.Name,
		Service:      yt.ReadableName,
		Filename:     id + ".track",
		ThumbnailURL: thumbnail,
		Duration:     duration,
		Playlist:     nil,
	}, nil
}
<commit_msg>Fix https://github.com/matthieugrieger/mumbledj/issues/180: Panic on playlist with private video<commit_after>/*
 * MumbleDJ
 * By Matthieu Grieger
 * services/youtube.go
 * Copyright (c) 2016 Matthieu Grieger (MIT License)
 */

package services

import (
	"errors"
	"fmt"
	"math"
	"net/http"
	"regexp"

	"github.com/ChannelMeter/iso8601duration"
	"github.com/antonholmquist/jason"
	"github.com/layeh/gumble/gumble"
	"github.com/matthieugrieger/mumbledj/bot"
	"github.com/matthieugrieger/mumbledj/interfaces"
	"github.com/spf13/viper"
)

// YouTube is a wrapper around the YouTube Data API.
// https://developers.google.com/youtube/v3/docs/
type YouTube struct {
	*GenericService
}

// NewYouTubeService returns an initialized YouTube service object.
func NewYouTubeService() *YouTube {
	return &YouTube{
		&GenericService{
			ReadableName: "YouTube",
			Format:       "bestaudio",
			TrackRegex: []*regexp.Regexp{
				regexp.MustCompile(`https?:\/\/www.youtube.com\/watch\?v=(?P<id>[\w-]+)(?P<timestamp>\&t=\d*m?\d*s?)?`),
				regexp.MustCompile(`https?:\/\/youtube.com\/watch\?v=(?P<id>[\w-]+)(?P<timestamp>\&t=\d*m?\d*s?)?`),
				regexp.MustCompile(`https?:\/\/youtu.be\/(?P<id>[\w-]+)(?P<timestamp>\?t=\d*m?\d*s?)?`),
				regexp.MustCompile(`https?:\/\/youtube.com\/v\/(?P<id>[\w-]+)(?P<timestamp>\?t=\d*m?\d*s?)?`),
				regexp.MustCompile(`https?:\/\/www.youtube.com\/v\/(?P<id>[\w-]+)(?P<timestamp>\?t=\d*m?\d*s?)?`),
			},
			PlaylistRegex: []*regexp.Regexp{
				regexp.MustCompile(`https?:\/\/www\.youtube\.com\/playlist\?list=(?P<id>[\w-]+)`),
			},
		},
	}
}

// CheckAPIKey performs a test API call with the API key
// provided in the configuration file to determine if the
// service should be enabled.
func (yt *YouTube) CheckAPIKey() error {
	var (
		response *http.Response
		v        *jason.Object
		err      error
	)

	if viper.GetString("api_keys.youtube") == "" {
		return errors.New("No YouTube API key has been provided")
	}
	url := "https://www.googleapis.com/youtube/v3/videos?part=snippet&id=KQY9zrjPBjo&key=%s"
	response, err = http.Get(fmt.Sprintf(url, viper.GetString("api_keys.youtube")))
	defer response.Body.Close()
	if err != nil {
		return err
	}

	if v, err = jason.NewObjectFromReader(response.Body); err != nil {
		return err
	}

	if v, err = v.GetObject("error"); err == nil {
		message, _ := v.GetString("message")
		code, _ := v.GetInt64("code")
		errArray, _ := v.GetObjectArray("errors")
		reason, _ := errArray[0].GetString("reason")

		return fmt.Errorf("%d: %s (reason: %s)", code, message, reason)
	}
	return nil
}

// GetTracks uses the passed URL to find and return
// tracks associated with the URL. An error is returned
// if any error occurs during the API call.
func (yt *YouTube) GetTracks(url string, submitter *gumble.User) ([]interfaces.Track, error) {
	var (
		playlistURL      string
		playlistItemsURL string
		id               string
		err              error
		resp             *http.Response
		v                *jason.Object
		track            bot.Track
		tracks           []interfaces.Track
	)

	playlistURL = "https://www.googleapis.com/youtube/v3/playlists?part=snippet&id=%s&key=%s"
	playlistItemsURL = "https://www.googleapis.com/youtube/v3/playlistItems?part=snippet,contentDetails&playlistId=%s&maxResults=%d&key=%s&pageToken=%s"
	id, err = yt.getID(url)
	if err != nil {
		return nil, err
	}

	if yt.isPlaylist(url) {
		resp, err = http.Get(fmt.Sprintf(playlistURL, id, viper.GetString("api_keys.youtube")))
		defer resp.Body.Close()
		if err != nil {
			return nil, err
		}

		v, err = jason.NewObjectFromReader(resp.Body)
		if err != nil {
			return nil, err
		}

		items, _ := v.GetObjectArray("items")
		item := items[0]

		title, _ := item.GetString("snippet", "title")

		playlist := &bot.Playlist{
			ID:        id,
			Title:     title,
			Submitter: submitter.Name,
			Service:   yt.ReadableName,
		}

		maxItems := math.MaxInt32
		if viper.GetInt("queue.max_tracks_per_playlist") > 0 {
			maxItems = viper.GetInt("queue.max_tracks_per_playlist")
		}

		// YouTube playlist searches return a max of 50 results per page
		maxResults := 50
		if maxResults > maxItems {
			maxResults = maxItems
		}

		pageToken := ""
		for len(tracks) < maxItems {
			curResp, curErr := http.Get(fmt.Sprintf(playlistItemsURL, id, maxResults, viper.GetString("api_keys.youtube"), pageToken))
			defer curResp.Body.Close()
			if curErr != nil {
				// An error occurred, simply skip this track.
				continue
			}

			v, err = jason.NewObjectFromReader(curResp.Body)
			if err != nil {
				// An error occurred, simply skip this track.
				continue
			}

			curTracks, _ := v.GetObjectArray("items")
			for _, track := range curTracks {
				videoID, _ := track.GetString("snippet", "resourceId", "videoId")

				// Unfortunately we have to execute another API call for each video as the YouTube API does not
				// return video durations from the playlistItems endpoint...
				newTrack, _ := yt.getTrack(videoID, submitter)
				newTrack.Playlist = playlist
				tracks = append(tracks, newTrack)

				if len(tracks) >= maxItems {
					break
				}
			}

			pageToken, _ = v.GetString("nextPageToken")
			if pageToken == "" {
				break
			}
		}

		if len(tracks) == 0 {
			return nil, errors.New("Invalid playlist. No tracks were added")
		}
		return tracks, nil
	}

	track, err = yt.getTrack(id, submitter)
	if err != nil {
		return nil, err
	}
	tracks = append(tracks, track)
	return tracks, nil
}

func (yt *YouTube) getTrack(id string, submitter *gumble.User) (bot.Track, error) {
	var (
		resp *http.Response
		err  error
		v    *jason.Object
	)

	videoURL := "https://www.googleapis.com/youtube/v3/videos?part=snippet,contentDetails&id=%s&key=%s"
	resp, err = http.Get(fmt.Sprintf(videoURL, id, viper.GetString("api_keys.youtube")))
	defer resp.Body.Close()
	if err != nil {
		return bot.Track{}, err
	}

	v, err = jason.NewObjectFromReader(resp.Body)
	if err != nil {
		return bot.Track{}, err
	}
	items, _ := v.GetObjectArray("items")
	if len(items) == 0 {
		return bot.Track{}, errors.New("This YouTube video is private")
	}
	item := items[0]
	title, _ := item.GetString("snippet", "title")
	thumbnail, _ := item.GetString("snippet", "thumbnails", "high", "url")
	author, _ := item.GetString("snippet", "channelTitle")
	durationString, _ := item.GetString("contentDetails", "duration")
	durationConverted, _ := duration.FromString(durationString)
	duration := durationConverted.ToDuration()

	return bot.Track{
		ID:           id,
		URL:          "https://youtube.com/watch?v=" + id,
		Title:        title,
		Author:       author,
		Submitter:    submitter.Name,
		Service:      yt.ReadableName,
		Filename:     id + ".track",
		ThumbnailURL: thumbnail,
		Duration:     duration,
		Playlist:     nil,
	}, nil
}
<|endoftext|>


<commit_before>package slack

/* util has most of the struct, type, and const definitions, as well as
most of the internal methods. */

import (
	"fmt"
	"log"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/nlopes/slack"
	"github.com/parsley42/gobot/bot"
)

const optimeout = 1 * time.Minute

type slackConnector struct {
	api          *slack.Client
	conn         *slack.RTM
	botName      string            // human-readable name of bot
	botID        string            // slack internal bot ID
	bot.Handler                    // bot interface for handling messages
	sync.RWMutex                   // shared mutex for locking connector data structures
	channelToID  map[string]string // map from channel names to channel IDs
	idToChannel  map[string]string // map from channel ID to channel name
	userToID     map[string]string // map from user names to user IDs
	idToUser     map[string]string // map from user name to user ID
	userIDToIM   map[string]string // map from user ID to IM channel ID
	imToUser     map[string]string // map from IM channel ID to user name
	level        bot.LogLevel      // current log level
}

func (s *slackConnector) addMessageMentions(msg string) string {
	re := regexp.MustCompile(`(?:^| )(@[0-9a-z]{1,21})\b`)
	mentions := re.FindAllStringSubmatch(msg, -1)
	if len(mentions) == 0 {
		return msg
	}
	mset := make(map[string]bool)
	for _, mention := range mentions {
		mset[mention[1]] = true
	}
	for mention, _ := range mset {
		userName := mention[1:]
		replace, ok := s.userID(userName)
		if !ok {
			// it's ok to have @notactuallyauser
			continue
		}
		msg = strings.Replace(msg, mention, "<@"+replace+">", -1)
	}
	return msg
}

// Examine incoming messages and route them to the appropriate bot
// method.
func (s *slackConnector) processMessage(msg *slack.MessageEvent) {
	s.log(bot.Trace, fmt.Sprintf("Message received: %v\n", msg))
	re := regexp.MustCompile(`<@U[A-Z0-9]{8}>`) // match a @user mention
	text := msg.Msg.Text
	chanID := msg.Msg.Channel
	mentions := re.FindAllString(text, -1)
	if len(mentions) != 0 {
		mset := make(map[string]bool)
		for _, mention := range mentions {
			mset[mention] = true
		}
		for mention, _ := range mset {
			mID := mention[2:11]
			replace, ok := s.userName(mID)
			if !ok {
				s.log(bot.Warn, "Couldn't find username for mentioned", mID)
				continue
			}
			text = strings.Replace(text, mention, "@"+replace, -1)
		}
	}
	switch chanID[:1] {
	case "D":
		userName, ok := s.imUser(chanID)
		if !ok {
			s.log(bot.Warn, "Couldn't find user name for IM", chanID)
			s.DirectMsg(chanID, text)
			return
		}
		s.DirectMsg(userName, text)
	case "C":
		channelName, ok := s.channelName(chanID)
		if !ok {
			s.log(bot.Warn, "Coudln't find channel name for ID", chanID)
			s.ChannelMsg(chanID, text)
			return
		}
		s.ChannelMsg(channelName, text)
	}
}

// log logs messages whenever the connector log level is
// less than the given level
func (s *slackConnector) log(l bot.LogLevel, v ...interface{}) {
	if l >= s.level {
		var prefix string
		switch l {
		case bot.Trace:
			prefix = "Trace:"
		case bot.Debug:
			prefix = "Debug:"
		case bot.Info:
			prefix = "Info"
		case bot.Warn:
			prefix = "Warning:"
		case bot.Error:
			prefix = "Error"
		}
		log.Println(prefix, v)
	}
}

// update maps is called whenever there are any changes
// to users or channels, so that plugins can use
// human-readable names for users and channels.
func (s *slackConnector) updateMaps() {
	s.log(bot.Trace, "Updating maps")
	deadline := time.Now().Add(optimeout)
	var (
		err        error
		userlist   []slack.User
		userIMlist []slack.IM
		chanlist   []slack.Channel
	)

	for tries := uint(0); time.Now().Before(deadline); tries++ {
		userlist, err = s.api.GetUsers()
		if err == nil {
			break
		}
	}
	if err != nil {
		log.Fatalf("Protocol timeout updating users: %v\n", err)
	}
	userMap := make(map[string]string)
	userIDMap := make(map[string]string)
	for _, user := range userlist {
		s.log(bot.Trace, "Mapping user name", user.Name, "to", user.ID)
		userMap[user.Name] = user.ID
		userIDMap[user.ID] = user.Name
	}

	for tries := uint(0); time.Now().Before(deadline); tries++ {
		userIMlist, err = s.api.GetIMChannels()
		if err == nil {
			break
		}
	}
	if err != nil {
		log.Fatalf("Protocol timeout updating IMchannels: %v\n", err)
	}
	userIMMap := make(map[string]string)
	userNameMap := make(map[string]string)
	for _, userIM := range userIMlist {
		s.log(bot.Trace, "Mapping user ID", userIM.User, "to IM channel", userIM.ID)
		userIMMap[userIM.User] = userIM.ID
		s.log(bot.Trace, "Mapping IM channel", userIM.ID, "to user name", userIDMap[userIM.User])
		userNameMap[userIM.ID] = userIDMap[userIM.User]
	}

	for tries := uint(0); time.Now().Before(deadline); tries++ {
		chanlist, err = s.api.GetChannels(true)
		if err == nil {
			break
		}
	}
	if err != nil {
		log.Fatalf("Protocol timeout updating channels: %v\n", err)
	}
	chanMap := make(map[string]string)
	chanIDMap := make(map[string]string)
	for _, channel := range chanlist {
		s.log(bot.Trace, "Mapping channel name", channel.Name, "to", channel.ID)
		chanMap[channel.Name] = channel.ID
		chanIDMap[channel.ID] = channel.Name
	}

	s.Lock()
	s.userToID = userMap
	s.idToUser = userIDMap
	s.userIDToIM = userIMMap
	s.channelToID = chanMap
	s.idToChannel = chanIDMap
	s.imToUser = userNameMap
	s.Unlock()
	s.log(bot.Info, "Users updated")
}

func (s *slackConnector) userID(u string) (i string, ok bool) {
	s.RLock()
	i, ok = s.userToID[u]
	s.RUnlock()
	return i, ok
}

func (s *slackConnector) userName(i string) (u string, ok bool) {
	s.RLock()
	u, ok = s.idToUser[i]
	s.RUnlock()
	return u, ok
}

func (s *slackConnector) userIMID(u string) (i string, ok bool) {
	s.RLock()
	i, ok = s.userIDToIM[u]
	s.RUnlock()
	return i, ok
}

func (s *slackConnector) chanID(c string) (i string, ok bool) {
	s.RLock()
	i, ok = s.channelToID[c]
	s.RUnlock()
	return i, ok
}

func (s *slackConnector) channelName(i string) (c string, ok bool) {
	s.RLock()
	c, ok = s.idToChannel[i]
	s.RUnlock()
	return c, ok
}

func (s *slackConnector) imUser(c string) (u string, ok bool) {
	s.RLock()
	u, ok = s.imToUser[c]
	s.RUnlock()
	return u, ok
}
<commit_msg>More succint outgoing mention replacement<commit_after>package slack

/* util has most of the struct, type, and const definitions, as well as
most of the internal methods. */

import (
	"fmt"
	"log"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/nlopes/slack"
	"github.com/parsley42/gobot/bot"
)

const optimeout = 1 * time.Minute

type slackConnector struct {
	api          *slack.Client
	conn         *slack.RTM
	botName      string            // human-readable name of bot
	botID        string            // slack internal bot ID
	bot.Handler                    // bot interface for handling messages
	sync.RWMutex                   // shared mutex for locking connector data structures
	channelToID  map[string]string // map from channel names to channel IDs
	idToChannel  map[string]string // map from channel ID to channel name
	userToID     map[string]string // map from user names to user IDs
	idToUser     map[string]string // map from user name to user ID
	userIDToIM   map[string]string // map from user ID to IM channel ID
	imToUser     map[string]string // map from IM channel ID to user name
	level        bot.LogLevel      // current log level
}

func (s *slackConnector) addMessageMentions(msg string) string {
	re := regexp.MustCompile(`@[0-9a-z]{1,21}\b`)
	return re.ReplaceAllStringFunc(msg, func(str string) string {
		replace, ok := s.userID(str[1:])
		if ok {
			return "<@" + replace + ">"
		}
		return str
	})
}

// Examine incoming messages and route them to the appropriate bot
// method.
func (s *slackConnector) processMessage(msg *slack.MessageEvent) {
	s.log(bot.Trace, fmt.Sprintf("Message received: %v\n", msg))
	re := regexp.MustCompile(`<@U[A-Z0-9]{8}>`) // match a @user mention
	text := msg.Msg.Text
	chanID := msg.Msg.Channel
	mentions := re.FindAllString(text, -1)
	if len(mentions) != 0 {
		mset := make(map[string]bool)
		for _, mention := range mentions {
			mset[mention] = true
		}
		for mention, _ := range mset {
			mID := mention[2:11]
			replace, ok := s.userName(mID)
			if !ok {
				s.log(bot.Warn, "Couldn't find username for mentioned", mID)
				continue
			}
			text = strings.Replace(text, mention, "@"+replace, -1)
		}
	}
	switch chanID[:1] {
	case "D":
		userName, ok := s.imUser(chanID)
		if !ok {
			s.log(bot.Warn, "Couldn't find user name for IM", chanID)
			s.DirectMsg(chanID, text)
			return
		}
		s.DirectMsg(userName, text)
	case "C":
		channelName, ok := s.channelName(chanID)
		if !ok {
			s.log(bot.Warn, "Coudln't find channel name for ID", chanID)
			s.ChannelMsg(chanID, text)
			return
		}
		s.ChannelMsg(channelName, text)
	}
}

// log logs messages whenever the connector log level is
// less than the given level
func (s *slackConnector) log(l bot.LogLevel, v ...interface{}) {
	if l >= s.level {
		var prefix string
		switch l {
		case bot.Trace:
			prefix = "Trace:"
		case bot.Debug:
			prefix = "Debug:"
		case bot.Info:
			prefix = "Info"
		case bot.Warn:
			prefix = "Warning:"
		case bot.Error:
			prefix = "Error"
		}
		log.Println(prefix, v)
	}
}

// update maps is called whenever there are any changes
// to users or channels, so that plugins can use
// human-readable names for users and channels.
func (s *slackConnector) updateMaps() {
	s.log(bot.Trace, "Updating maps")
	deadline := time.Now().Add(optimeout)
	var (
		err        error
		userlist   []slack.User
		userIMlist []slack.IM
		chanlist   []slack.Channel
	)

	for tries := uint(0); time.Now().Before(deadline); tries++ {
		userlist, err = s.api.GetUsers()
		if err == nil {
			break
		}
	}
	if err != nil {
		log.Fatalf("Protocol timeout updating users: %v\n", err)
	}
	userMap := make(map[string]string)
	userIDMap := make(map[string]string)
	for _, user := range userlist {
		s.log(bot.Trace, "Mapping user name", user.Name, "to", user.ID)
		userMap[user.Name] = user.ID
		userIDMap[user.ID] = user.Name
	}

	for tries := uint(0); time.Now().Before(deadline); tries++ {
		userIMlist, err = s.api.GetIMChannels()
		if err == nil {
			break
		}
	}
	if err != nil {
		log.Fatalf("Protocol timeout updating IMchannels: %v\n", err)
	}
	userIMMap := make(map[string]string)
	userNameMap := make(map[string]string)
	for _, userIM := range userIMlist {
		s.log(bot.Trace, "Mapping user ID", userIM.User, "to IM channel", userIM.ID)
		userIMMap[userIM.User] = userIM.ID
		s.log(bot.Trace, "Mapping IM channel", userIM.ID, "to user name", userIDMap[userIM.User])
		userNameMap[userIM.ID] = userIDMap[userIM.User]
	}

	for tries := uint(0); time.Now().Before(deadline); tries++ {
		chanlist, err = s.api.GetChannels(true)
		if err == nil {
			break
		}
	}
	if err != nil {
		log.Fatalf("Protocol timeout updating channels: %v\n", err)
	}
	chanMap := make(map[string]string)
	chanIDMap := make(map[string]string)
	for _, channel := range chanlist {
		s.log(bot.Trace, "Mapping channel name", channel.Name, "to", channel.ID)
		chanMap[channel.Name] = channel.ID
		chanIDMap[channel.ID] = channel.Name
	}

	s.Lock()
	s.userToID = userMap
	s.idToUser = userIDMap
	s.userIDToIM = userIMMap
	s.channelToID = chanMap
	s.idToChannel = chanIDMap
	s.imToUser = userNameMap
	s.Unlock()
	s.log(bot.Info, "Users updated")
}

func (s *slackConnector) userID(u string) (i string, ok bool) {
	s.RLock()
	i, ok = s.userToID[u]
	s.RUnlock()
	return i, ok
}

func (s *slackConnector) userName(i string) (u string, ok bool) {
	s.RLock()
	u, ok = s.idToUser[i]
	s.RUnlock()
	return u, ok
}

func (s *slackConnector) userIMID(u string) (i string, ok bool) {
	s.RLock()
	i, ok = s.userIDToIM[u]
	s.RUnlock()
	return i, ok
}

func (s *slackConnector) chanID(c string) (i string, ok bool) {
	s.RLock()
	i, ok = s.channelToID[c]
	s.RUnlock()
	return i, ok
}

func (s *slackConnector) channelName(i string) (c string, ok bool) {
	s.RLock()
	c, ok = s.idToChannel[i]
	s.RUnlock()
	return c, ok
}

func (s *slackConnector) imUser(c string) (u string, ok bool) {
	s.RLock()
	u, ok = s.imToUser[c]
	s.RUnlock()
	return u, ok
}
<|endoftext|>


<commit_before>package main

import (
	"path"
	"runtime"

	"go.skia.org/infra/go/auth"
	"go.skia.org/infra/go/gce"
	"go.skia.org/infra/go/gce/server"
)

func TaskSchedulerBase(name, ipAddress string) *gce.Instance {
	vm := server.Server20170928(name)
	vm.DataDisks[0].SizeGb = 200
	vm.DataDisks[0].Type = gce.DISK_TYPE_PERSISTENT_SSD
	vm.ExternalIpAddress = ipAddress
	vm.Metadata["owner_primary"] = "borenet"
	vm.Metadata["owner_secondary"] = "benjaminwagner"
	vm.Scopes = append(vm.Scopes,
		auth.SCOPE_GERRIT,
	)

	_, filename, _, _ := runtime.Caller(0)
	dir := path.Dir(filename)
	vm.StartupScript = path.Join(dir, "startup-script.sh")
	return vm
}

func TaskSchedulerProd() *gce.Instance {
	vm := TaskSchedulerBase("skia-task-scheduler", "35.202.175.145" /* Whitelisted in swarming, isolate and buildbucket servers */)
	vm.DataDisks[0].SizeGb = 500
	vm.ServiceAccount = "task-scheduler@skia-buildbots.google.com.iam.gserviceaccount.com"
	return vm
}

func TaskSchedulerInternal() *gce.Instance {
	vm := TaskSchedulerBase("skia-task-scheduler-internal", "35.184.167.88" /* Whitelisted in swarming, isolate and buildbucket servers */)
	vm.ServiceAccount = "task-scheduler-internal@skia-buildbots.google.com.iam.gserviceaccount.com"
	return vm
}

func TaskSchedulerStaging() *gce.Instance {
	vm := TaskSchedulerBase("skia-task-scheduler-staging", "")
	vm.ServiceAccount = "task-scheduler@skia-buildbots.google.com.iam.gserviceaccount.com"
	return vm
}

func main() {
	server.Main(gce.ZONE_DEFAULT, map[string]*gce.Instance{
		"prod":     TaskSchedulerProd(),
		"internal": TaskSchedulerInternal(),
		"staging":  TaskSchedulerStaging(),
	})
}
<commit_msg>[task scheduler] Use static IP for staging instance<commit_after>package main

import (
	"path"
	"runtime"

	"go.skia.org/infra/go/auth"
	"go.skia.org/infra/go/gce"
	"go.skia.org/infra/go/gce/server"
)

func TaskSchedulerBase(name, ipAddress string) *gce.Instance {
	vm := server.Server20170928(name)
	vm.DataDisks[0].SizeGb = 200
	vm.DataDisks[0].Type = gce.DISK_TYPE_PERSISTENT_SSD
	vm.ExternalIpAddress = ipAddress
	vm.Metadata["owner_primary"] = "borenet"
	vm.Metadata["owner_secondary"] = "benjaminwagner"
	vm.Scopes = append(vm.Scopes,
		auth.SCOPE_GERRIT,
	)

	_, filename, _, _ := runtime.Caller(0)
	dir := path.Dir(filename)
	vm.StartupScript = path.Join(dir, "startup-script.sh")
	return vm
}

func TaskSchedulerProd() *gce.Instance {
	vm := TaskSchedulerBase("skia-task-scheduler", "35.202.175.145" /* Whitelisted in swarming, isolate and buildbucket servers */)
	vm.DataDisks[0].SizeGb = 500
	vm.ServiceAccount = "task-scheduler@skia-buildbots.google.com.iam.gserviceaccount.com"
	return vm
}

func TaskSchedulerInternal() *gce.Instance {
	vm := TaskSchedulerBase("skia-task-scheduler-internal", "35.184.167.88" /* Whitelisted in swarming, isolate and buildbucket servers */)
	vm.ServiceAccount = "task-scheduler-internal@skia-buildbots.google.com.iam.gserviceaccount.com"
	return vm
}

func TaskSchedulerStaging() *gce.Instance {
	vm := TaskSchedulerBase("skia-task-scheduler-staging", "35.239.145.127" /* Whitelisted in swarming, isolate and buildbucket servers */)
	vm.ServiceAccount = "task-scheduler@skia-buildbots.google.com.iam.gserviceaccount.com"
	return vm
}

func main() {
	server.Main(gce.ZONE_DEFAULT, map[string]*gce.Instance{
		"prod":     TaskSchedulerProd(),
		"internal": TaskSchedulerInternal(),
		"staging":  TaskSchedulerStaging(),
	})
}
<|endoftext|>


<commit_before>package sftpd

import (
	"io"
	"log"
	"os"
	"github.com/sandreas/sftp"
)

type vfs struct {
	pathMap PathMapper
}

func VfsHandler(mapper *PathMapper) sftp.Handlers {
	virtualFileSystem := &vfs{
		pathMap: *mapper,
	}
	return sftp.Handlers{
		FileGet: virtualFileSystem,
		FilePut: virtualFileSystem,
		FileCmd: virtualFileSystem,
		FileInfo: virtualFileSystem,
	}
}

func dumpSftpRequest(message string, r sftp.Request) {
	log.Println(message, "Filepath: ", r.Filepath, ", Target: ", r.Target, ", Method: ", r.Method)
}

func (fs *vfs) Fileread(r sftp.Request) (io.ReaderAt, error) {
	dumpSftpRequest("Fileread: ", r)

	filePath, err := fs.pathMap.PathTo(r.Filepath)
	if err == nil {
		return os.Open(filePath)
	}
	return nil, err
}

func (fs *vfs) Filewrite(r sftp.Request) (io.WriterAt, error) {
	dumpSftpRequest("Filewrite (disabled): ", r)
	return nil, os.ErrInvalid
}

func (fs *vfs) Filecmd(r sftp.Request) error {
	dumpSftpRequest("Filecmd (disabled): ", r)
	return os.ErrInvalid
}

func (fs *vfs) Fileinfo(r sftp.Request) ([]os.FileInfo, error) {
	dumpSftpRequest("Fileinfo: ", r)
	switch r.Method {
	case "List":
		listing, ok := fs.pathMap.List(r.Filepath)
		if ! ok {
			return nil, os.ErrInvalid
		}

		statList := make([]os.FileInfo, len(listing))
		for i, fileName := range listing {
			stat, err := fs.pathMap.Stat(fileName)
			if err != nil {
				log.Println("Could not stat file", fileName, err)
				continue
			}

			statList[i] = stat
			log.Println("Stat for file "+fileName+": isDir=>", stat.IsDir(), "size=>", stat.Size())
		}
		return statList, nil
	case "Stat":
		stat, err := fs.pathMap.Stat(r.Filepath)
		if err != nil {
			return nil, err
		}
		return []os.FileInfo{stat}, nil
	}
	return nil, os.ErrInvalid
}
<commit_msg>Experiments with defer close for vfshandler<commit_after>package sftpd

import (
	"io"
	"log"
	"os"
	"github.com/sandreas/sftp"
)

type vfs struct {
	pathMap PathMapper
}

func VfsHandler(mapper *PathMapper) sftp.Handlers {
	virtualFileSystem := &vfs{
		pathMap: *mapper,
	}
	return sftp.Handlers{
		FileGet: virtualFileSystem,
		FilePut: virtualFileSystem,
		FileCmd: virtualFileSystem,
		FileInfo: virtualFileSystem,
	}
}

func dumpSftpRequest(message string, r sftp.Request) {
	log.Println(message, "Filepath: ", r.Filepath, ", Target: ", r.Target, ", Method: ", r.Method)
}

func (fs *vfs) Fileread(r sftp.Request) (io.ReaderAt, error) {
	dumpSftpRequest("Fileread: ", r)

	filePath, err := fs.pathMap.PathTo(r.Filepath)

	if err == nil {
		f, err := os.Open(filePath)
		if err != nil {
			defer f.Close()
		}
		return f, err
	}
	return nil, err
}

func (fs *vfs) Filewrite(r sftp.Request) (io.WriterAt, error) {
	dumpSftpRequest("Filewrite (disabled): ", r)
	return nil, os.ErrInvalid
}

func (fs *vfs) Filecmd(r sftp.Request) error {
	dumpSftpRequest("Filecmd (disabled): ", r)
	return os.ErrInvalid
}

func (fs *vfs) Fileinfo(r sftp.Request) ([]os.FileInfo, error) {
	dumpSftpRequest("Fileinfo: ", r)
	switch r.Method {
	case "List":
		listing, ok := fs.pathMap.List(r.Filepath)
		if ! ok {
			return nil, os.ErrInvalid
		}

		statList := make([]os.FileInfo, len(listing))
		for i, fileName := range listing {
			stat, err := fs.pathMap.Stat(fileName)
			if err != nil {
				log.Println("Could not stat file", fileName, err)
				continue
			}

			statList[i] = stat
			log.Println("Stat for file "+fileName+": isDir=>", stat.IsDir(), "size=>", stat.Size())
		}
		return statList, nil
	case "Stat":
		stat, err := fs.pathMap.Stat(r.Filepath)
		if err != nil {
			return nil, err
		}
		return []os.FileInfo{stat}, nil
	}
	return nil, os.ErrInvalid
}
<|endoftext|>