<commit_before>package storage

import (
	"context"
	"fmt"
	"io"
	"sync"
	"time"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/operation"
	"github.com/chrislusf/seaweedfs/weed/pb/master_pb"
	"github.com/chrislusf/seaweedfs/weed/pb/volume_server_pb"
	"github.com/chrislusf/seaweedfs/weed/storage/erasure_coding"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
	"github.com/klauspost/reedsolomon"
)

func (s *Store) CollectErasureCodingHeartbeat() *master_pb.Heartbeat {
	var ecShardMessages []*master_pb.VolumeEcShardInformationMessage
	for _, location := range s.Locations {
		location.ecVolumesLock.RLock()
		for _, ecShards := range location.ecVolumes {
			ecShardMessages = append(ecShardMessages, ecShards.ToVolumeEcShardInformationMessage()...)
		}
		location.ecVolumesLock.RUnlock()
	}

	return &master_pb.Heartbeat{
		EcShards: ecShardMessages,
	}

}

func (s *Store) MountEcShards(collection string, vid needle.VolumeId, shardId erasure_coding.ShardId) error {
	for _, location := range s.Locations {
		if err := location.LoadEcShard(collection, vid, shardId); err == nil {
			glog.V(0).Infof("MountEcShards %d.%d", vid, shardId)

			var shardBits erasure_coding.ShardBits

			s.NewEcShardsChan <- master_pb.VolumeEcShardInformationMessage{
				Id:          uint32(vid),
				Collection:  collection,
				EcIndexBits: uint32(shardBits.AddShardId(shardId)),
			}
			return nil
		}
	}

	return fmt.Errorf("MountEcShards %d.%d not found on disk", vid, shardId)
}

func (s *Store) UnmountEcShards(vid needle.VolumeId, shardId erasure_coding.ShardId) error {

	ecShard, found := s.findEcShard(vid, shardId)
	if !found {
		return nil
	}

	var shardBits erasure_coding.ShardBits
	message := master_pb.VolumeEcShardInformationMessage{
		Id:          uint32(vid),
		Collection:  ecShard.Collection,
		EcIndexBits: uint32(shardBits.AddShardId(shardId)),
	}

	for _, location := range s.Locations {
		if deleted := location.UnloadEcShard(vid, shardId); deleted {
			glog.V(0).Infof("UnmountEcShards %d.%d", vid, shardId)
			s.DeletedEcShardsChan <- message
			return nil
		}
	}

	return fmt.Errorf("UnmountEcShards %d.%d not found on disk", vid, shardId)
}

func (s *Store) findEcShard(vid needle.VolumeId, shardId erasure_coding.ShardId) (*erasure_coding.EcVolumeShard, bool) {
	for _, location := range s.Locations {
		if v, found := location.FindEcShard(vid, shardId); found {
			return v, found
		}
	}
	return nil, false
}

func (s *Store) FindEcVolume(vid needle.VolumeId) (*erasure_coding.EcVolume, bool) {
	for _, location := range s.Locations {
		if s, found := location.FindEcVolume(vid); found {
			return s, true
		}
	}
	return nil, false
}

func (s *Store) ReadEcShardNeedle(ctx context.Context, vid needle.VolumeId, n *needle.Needle) (int, error) {
	for _, location := range s.Locations {
		if localEcVolume, found := location.FindEcVolume(vid); found {

			// read the volume version
			for localEcVolume.Version == 0 {
				err := s.readEcVolumeVersion(ctx, vid, localEcVolume)
				time.Sleep(1357 * time.Millisecond)
				glog.V(0).Infof("ReadEcShardNeedle vid %d version:%v: %v", vid, localEcVolume.Version, err)
			}
			version := localEcVolume.Version

			offset, size, intervals, err := localEcVolume.LocateEcShardNeedle(n, version)
			if err != nil {
				return 0, err
			}

			glog.V(4).Infof("read ec volume %d offset %d size %d intervals:%+v", vid, offset.ToAcutalOffset(), size, intervals)

			if len(intervals) > 1 {
				glog.V(4).Infof("ReadEcShardNeedle needle id %s intervals:%+v", n.String(), intervals)
			}
			bytes, err := s.readEcShardIntervals(ctx, vid, localEcVolume, intervals)
			if err != nil {
				return 0, fmt.Errorf("ReadEcShardIntervals: %v", err)
			}

			err = n.ReadBytes(bytes, offset.ToAcutalOffset(), size, version)
			if err != nil {
				return 0, fmt.Errorf("readbytes: %v", err)
			}

			return len(bytes), nil
		}
	}
	return 0, fmt.Errorf("ec shard %d not found", vid)
}

func (s *Store) readEcVolumeVersion(ctx context.Context, vid needle.VolumeId, ecVolume *erasure_coding.EcVolume) (err error) {

	interval := erasure_coding.Interval{
		BlockIndex:          0,
		InnerBlockOffset:    0,
		Size:                _SuperBlockSize,
		IsLargeBlock:        true, // it could be large block, but ok in this place
		LargeBlockRowsCount: 0,
	}
	data, err := s.readEcShardIntervals(ctx, vid, ecVolume, []erasure_coding.Interval{interval})
	if err == nil {
		ecVolume.Version = needle.Version(data[0])
	}
	return
}

func (s *Store) readEcShardIntervals(ctx context.Context, vid needle.VolumeId, ecVolume *erasure_coding.EcVolume, intervals []erasure_coding.Interval) (data []byte, err error) {

	if err = s.cachedLookupEcShardLocations(ctx, ecVolume); err != nil {
		return nil, fmt.Errorf("failed to locate shard via master grpc %s: %v", s.MasterAddress, err)
	}

	for i, interval := range intervals {
		if d, e := s.readOneEcShardInterval(ctx, ecVolume, interval); e != nil {
			return nil, e
		} else {
			if i == 0 {
				data = d
			} else {
				data = append(data, d...)
			}
		}
	}
	return
}

func (s *Store) readOneEcShardInterval(ctx context.Context, ecVolume *erasure_coding.EcVolume, interval erasure_coding.Interval) (data []byte, err error) {
	shardId, actualOffset := interval.ToShardIdAndOffset(erasure_coding.ErasureCodingLargeBlockSize, erasure_coding.ErasureCodingSmallBlockSize)
	data = make([]byte, interval.Size)
	if shard, found := ecVolume.FindEcVolumeShard(shardId); found {
		if _, err = shard.ReadAt(data, actualOffset); err != nil {
			glog.V(0).Infof("read local ec shard %d.%d: %v", ecVolume.VolumeId, shardId, err)
			return
		}
	} else {
		ecVolume.ShardLocationsLock.RLock()
		sourceDataNodes, hasShardIdLocation := ecVolume.ShardLocations[shardId]
		ecVolume.ShardLocationsLock.RUnlock()

		// try reading directly
		if hasShardIdLocation {
			_, err = s.readRemoteEcShardInterval(ctx, sourceDataNodes, ecVolume.VolumeId, shardId, data, actualOffset)
			if err == nil {
				return
			}
			glog.V(0).Infof("clearing ec shard %d.%d locations: %v", ecVolume.VolumeId, shardId, err)
			forgetShardId(ecVolume, shardId)
		}

		// try reading by recovering from other shards
		_, err = s.recoverOneRemoteEcShardInterval(ctx, ecVolume, shardId, data, actualOffset)
		if err == nil {
			return
		}
		glog.V(0).Infof("recover ec shard %d.%d : %v", ecVolume.VolumeId, shardId, err)
	}
	return
}

func forgetShardId(ecVolume *erasure_coding.EcVolume, shardId erasure_coding.ShardId) {
	// failed to access the source data nodes, clear it up
	ecVolume.ShardLocationsLock.Lock()
	delete(ecVolume.ShardLocations, shardId)
	ecVolume.ShardLocationsLock.Unlock()
}

func (s *Store) cachedLookupEcShardLocations(ctx context.Context, ecVolume *erasure_coding.EcVolume) (err error) {

	if ecVolume.ShardLocationsRefreshTime.Add(10 * time.Minute).After(time.Now()) {
		// still fresh
		return nil
	}

	glog.V(3).Infof("lookup and cache ec volume %d locations", ecVolume.VolumeId)

	err = operation.WithMasterServerClient(s.MasterAddress, s.grpcDialOption, func(masterClient master_pb.SeaweedClient) error {
		req := &master_pb.LookupEcVolumeRequest{
			VolumeId: uint32(ecVolume.VolumeId),
		}
		resp, err := masterClient.LookupEcVolume(ctx, req)
		if err != nil {
			return fmt.Errorf("lookup ec volume %d: %v", ecVolume.VolumeId, err)
		}
		if len(resp.ShardIdLocations) < erasure_coding.DataShardsCount {
			return fmt.Errorf("only %d shards found but %d required", len(resp.ShardIdLocations), erasure_coding.DataShardsCount)
		}

		ecVolume.ShardLocationsLock.Lock()
		for _, shardIdLocations := range resp.ShardIdLocations {
			shardId := erasure_coding.ShardId(shardIdLocations.ShardId)
			delete(ecVolume.ShardLocations, shardId)
			for _, loc := range shardIdLocations.Locations {
				ecVolume.ShardLocations[shardId] = append(ecVolume.ShardLocations[shardId], loc.Url)
			}
		}
		ecVolume.ShardLocationsRefreshTime = time.Now()
		ecVolume.ShardLocationsLock.Unlock()

		return nil
	})
	return
}

func (s *Store) readRemoteEcShardInterval(ctx context.Context, sourceDataNodes []string, vid needle.VolumeId, shardId erasure_coding.ShardId, buf []byte, offset int64) (n int, err error) {

	if len(sourceDataNodes) == 0 {
		return 0, fmt.Errorf("failed to find ec shard %d.%d", vid, shardId)
	}

	for _, sourceDataNode := range sourceDataNodes {
		glog.V(4).Infof("read remote ec shard %d.%d from %s", vid, shardId, sourceDataNode)
		n, err = s.doReadRemoteEcShardInterval(ctx, sourceDataNode, vid, shardId, buf, offset)
		if err == nil {
			return
		}
		glog.V(1).Infof("read remote ec shard %d.%d from %s: %v", vid, shardId, sourceDataNode, err)
	}

	return
}

func (s *Store) doReadRemoteEcShardInterval(ctx context.Context, sourceDataNode string, vid needle.VolumeId, shardId erasure_coding.ShardId, buf []byte, offset int64) (n int, err error) {

	err = operation.WithVolumeServerClient(sourceDataNode, s.grpcDialOption, func(client volume_server_pb.VolumeServerClient) error {

		// copy data slice
		shardReadClient, err := client.VolumeEcShardRead(ctx, &volume_server_pb.VolumeEcShardReadRequest{
			VolumeId: uint32(vid),
			ShardId:  uint32(shardId),
			Offset:   offset,
			Size:     int64(len(buf)),
		})
		if err != nil {
			return fmt.Errorf("failed to start reading ec shard %d.%d from %s: %v", vid, shardId, sourceDataNode, err)
		}

		for {
			resp, receiveErr := shardReadClient.Recv()
			if receiveErr == io.EOF {
				break
			}
			if receiveErr != nil {
				return fmt.Errorf("receiving ec shard %d.%d from %s: %v", vid, shardId, sourceDataNode, err)
			}
			copy(buf[n:n+len(resp.Data)], resp.Data)
			n += len(resp.Data)
		}

		return nil
	})
	if err != nil {
		return 0, fmt.Errorf("read ec shard %d.%d from %s: %v", vid, shardId, sourceDataNode, err)
	}

	return
}

func (s *Store) recoverOneRemoteEcShardInterval(ctx context.Context, ecVolume *erasure_coding.EcVolume, shardIdToRecover erasure_coding.ShardId, buf []byte, offset int64) (n int, err error) {
	glog.V(4).Infof("recover ec shard %d.%d from other locations", ecVolume.VolumeId, shardIdToRecover)

	enc, err := reedsolomon.New(erasure_coding.DataShardsCount, erasure_coding.ParityShardsCount)
	if err != nil {
		return 0, fmt.Errorf("failed to create encoder: %v", err)
	}

	bufs := make([][]byte, erasure_coding.TotalShardsCount)

	var wg sync.WaitGroup
	ecVolume.ShardLocationsLock.RLock()
	for shardId, locations := range ecVolume.ShardLocations {

		// skip currnent shard or empty shard
		if shardId == shardIdToRecover {
			continue
		}
		if len(locations) == 0 {
			glog.V(3).Infof("readRemoteEcShardInterval missing %d.%d from %+v", ecVolume.VolumeId, shardId, locations)
			continue
		}

		// read from remote locations
		wg.Add(1)
		go func(shardId erasure_coding.ShardId, locations []string) {
			defer wg.Done()
			data := make([]byte, len(buf))
			nRead, readErr := s.readRemoteEcShardInterval(ctx, locations, ecVolume.VolumeId, shardId, data, offset)
			if readErr != nil {
				glog.V(3).Infof("recover: readRemoteEcShardInterval %d.%d %d bytes from %+v: %v", ecVolume.VolumeId, shardId, nRead, locations, readErr)
				forgetShardId(ecVolume, shardId)
			}
			if nRead == len(buf) {
				bufs[shardId] = data
			}
		}(shardId, locations)
	}
	ecVolume.ShardLocationsLock.RUnlock()

	wg.Wait()

	if err = enc.ReconstructData(bufs); err != nil {
		glog.V(3).Infof("recovered ec shard %d.%d failed: %v", ecVolume.VolumeId, shardIdToRecover, err)
		return 0, err
	}
	glog.V(4).Infof("recovered ec shard %d.%d from other locations", ecVolume.VolumeId, shardIdToRecover)

	copy(buf, bufs[shardIdToRecover])

	return len(buf), nil
}
<commit_msg>conditionally fresh the shard locations<commit_after>package storage

import (
	"context"
	"fmt"
	"io"
	"sync"
	"time"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/operation"
	"github.com/chrislusf/seaweedfs/weed/pb/master_pb"
	"github.com/chrislusf/seaweedfs/weed/pb/volume_server_pb"
	"github.com/chrislusf/seaweedfs/weed/storage/erasure_coding"
	"github.com/chrislusf/seaweedfs/weed/storage/needle"
	"github.com/klauspost/reedsolomon"
)

func (s *Store) CollectErasureCodingHeartbeat() *master_pb.Heartbeat {
	var ecShardMessages []*master_pb.VolumeEcShardInformationMessage
	for _, location := range s.Locations {
		location.ecVolumesLock.RLock()
		for _, ecShards := range location.ecVolumes {
			ecShardMessages = append(ecShardMessages, ecShards.ToVolumeEcShardInformationMessage()...)
		}
		location.ecVolumesLock.RUnlock()
	}

	return &master_pb.Heartbeat{
		EcShards: ecShardMessages,
	}

}

func (s *Store) MountEcShards(collection string, vid needle.VolumeId, shardId erasure_coding.ShardId) error {
	for _, location := range s.Locations {
		if err := location.LoadEcShard(collection, vid, shardId); err == nil {
			glog.V(0).Infof("MountEcShards %d.%d", vid, shardId)

			var shardBits erasure_coding.ShardBits

			s.NewEcShardsChan <- master_pb.VolumeEcShardInformationMessage{
				Id:          uint32(vid),
				Collection:  collection,
				EcIndexBits: uint32(shardBits.AddShardId(shardId)),
			}
			return nil
		}
	}

	return fmt.Errorf("MountEcShards %d.%d not found on disk", vid, shardId)
}

func (s *Store) UnmountEcShards(vid needle.VolumeId, shardId erasure_coding.ShardId) error {

	ecShard, found := s.findEcShard(vid, shardId)
	if !found {
		return nil
	}

	var shardBits erasure_coding.ShardBits
	message := master_pb.VolumeEcShardInformationMessage{
		Id:          uint32(vid),
		Collection:  ecShard.Collection,
		EcIndexBits: uint32(shardBits.AddShardId(shardId)),
	}

	for _, location := range s.Locations {
		if deleted := location.UnloadEcShard(vid, shardId); deleted {
			glog.V(0).Infof("UnmountEcShards %d.%d", vid, shardId)
			s.DeletedEcShardsChan <- message
			return nil
		}
	}

	return fmt.Errorf("UnmountEcShards %d.%d not found on disk", vid, shardId)
}

func (s *Store) findEcShard(vid needle.VolumeId, shardId erasure_coding.ShardId) (*erasure_coding.EcVolumeShard, bool) {
	for _, location := range s.Locations {
		if v, found := location.FindEcShard(vid, shardId); found {
			return v, found
		}
	}
	return nil, false
}

func (s *Store) FindEcVolume(vid needle.VolumeId) (*erasure_coding.EcVolume, bool) {
	for _, location := range s.Locations {
		if s, found := location.FindEcVolume(vid); found {
			return s, true
		}
	}
	return nil, false
}

func (s *Store) ReadEcShardNeedle(ctx context.Context, vid needle.VolumeId, n *needle.Needle) (int, error) {
	for _, location := range s.Locations {
		if localEcVolume, found := location.FindEcVolume(vid); found {

			// read the volume version
			for localEcVolume.Version == 0 {
				err := s.readEcVolumeVersion(ctx, vid, localEcVolume)
				time.Sleep(1357 * time.Millisecond)
				glog.V(0).Infof("ReadEcShardNeedle vid %d version:%v: %v", vid, localEcVolume.Version, err)
			}
			version := localEcVolume.Version

			offset, size, intervals, err := localEcVolume.LocateEcShardNeedle(n, version)
			if err != nil {
				return 0, err
			}

			glog.V(4).Infof("read ec volume %d offset %d size %d intervals:%+v", vid, offset.ToAcutalOffset(), size, intervals)

			if len(intervals) > 1 {
				glog.V(4).Infof("ReadEcShardNeedle needle id %s intervals:%+v", n.String(), intervals)
			}
			bytes, err := s.readEcShardIntervals(ctx, vid, localEcVolume, intervals)
			if err != nil {
				return 0, fmt.Errorf("ReadEcShardIntervals: %v", err)
			}

			err = n.ReadBytes(bytes, offset.ToAcutalOffset(), size, version)
			if err != nil {
				return 0, fmt.Errorf("readbytes: %v", err)
			}

			return len(bytes), nil
		}
	}
	return 0, fmt.Errorf("ec shard %d not found", vid)
}

func (s *Store) readEcVolumeVersion(ctx context.Context, vid needle.VolumeId, ecVolume *erasure_coding.EcVolume) (err error) {

	interval := erasure_coding.Interval{
		BlockIndex:          0,
		InnerBlockOffset:    0,
		Size:                _SuperBlockSize,
		IsLargeBlock:        true, // it could be large block, but ok in this place
		LargeBlockRowsCount: 0,
	}
	data, err := s.readEcShardIntervals(ctx, vid, ecVolume, []erasure_coding.Interval{interval})
	if err == nil {
		ecVolume.Version = needle.Version(data[0])
	}
	return
}

func (s *Store) readEcShardIntervals(ctx context.Context, vid needle.VolumeId, ecVolume *erasure_coding.EcVolume, intervals []erasure_coding.Interval) (data []byte, err error) {

	if err = s.cachedLookupEcShardLocations(ctx, ecVolume); err != nil {
		return nil, fmt.Errorf("failed to locate shard via master grpc %s: %v", s.MasterAddress, err)
	}

	for i, interval := range intervals {
		if d, e := s.readOneEcShardInterval(ctx, ecVolume, interval); e != nil {
			return nil, e
		} else {
			if i == 0 {
				data = d
			} else {
				data = append(data, d...)
			}
		}
	}
	return
}

func (s *Store) readOneEcShardInterval(ctx context.Context, ecVolume *erasure_coding.EcVolume, interval erasure_coding.Interval) (data []byte, err error) {
	shardId, actualOffset := interval.ToShardIdAndOffset(erasure_coding.ErasureCodingLargeBlockSize, erasure_coding.ErasureCodingSmallBlockSize)
	data = make([]byte, interval.Size)
	if shard, found := ecVolume.FindEcVolumeShard(shardId); found {
		if _, err = shard.ReadAt(data, actualOffset); err != nil {
			glog.V(0).Infof("read local ec shard %d.%d: %v", ecVolume.VolumeId, shardId, err)
			return
		}
	} else {
		ecVolume.ShardLocationsLock.RLock()
		sourceDataNodes, hasShardIdLocation := ecVolume.ShardLocations[shardId]
		ecVolume.ShardLocationsLock.RUnlock()

		// try reading directly
		if hasShardIdLocation {
			_, err = s.readRemoteEcShardInterval(ctx, sourceDataNodes, ecVolume.VolumeId, shardId, data, actualOffset)
			if err == nil {
				return
			}
			glog.V(0).Infof("clearing ec shard %d.%d locations: %v", ecVolume.VolumeId, shardId, err)
			forgetShardId(ecVolume, shardId)
		}

		// try reading by recovering from other shards
		_, err = s.recoverOneRemoteEcShardInterval(ctx, ecVolume, shardId, data, actualOffset)
		if err == nil {
			return
		}
		glog.V(0).Infof("recover ec shard %d.%d : %v", ecVolume.VolumeId, shardId, err)
	}
	return
}

func forgetShardId(ecVolume *erasure_coding.EcVolume, shardId erasure_coding.ShardId) {
	// failed to access the source data nodes, clear it up
	ecVolume.ShardLocationsLock.Lock()
	delete(ecVolume.ShardLocations, shardId)
	ecVolume.ShardLocationsLock.Unlock()
}

func (s *Store) cachedLookupEcShardLocations(ctx context.Context, ecVolume *erasure_coding.EcVolume) (err error) {

	shardCount := len(ecVolume.ShardLocations)
	if shardCount < erasure_coding.DataShardsCount &&
		ecVolume.ShardLocationsRefreshTime.Add(11 * time.Second).After(time.Now()) ||
		shardCount == erasure_coding.TotalShardsCount &&
			ecVolume.ShardLocationsRefreshTime.Add(37 * time.Minute).After(time.Now()) ||
		shardCount >= erasure_coding.DataShardsCount &&
			ecVolume.ShardLocationsRefreshTime.Add(7 * time.Minute).After(time.Now()) {
		// still fresh
		return nil
	}

	glog.V(3).Infof("lookup and cache ec volume %d locations", ecVolume.VolumeId)

	err = operation.WithMasterServerClient(s.MasterAddress, s.grpcDialOption, func(masterClient master_pb.SeaweedClient) error {
		req := &master_pb.LookupEcVolumeRequest{
			VolumeId: uint32(ecVolume.VolumeId),
		}
		resp, err := masterClient.LookupEcVolume(ctx, req)
		if err != nil {
			return fmt.Errorf("lookup ec volume %d: %v", ecVolume.VolumeId, err)
		}
		if len(resp.ShardIdLocations) < erasure_coding.DataShardsCount {
			return fmt.Errorf("only %d shards found but %d required", len(resp.ShardIdLocations), erasure_coding.DataShardsCount)
		}

		ecVolume.ShardLocationsLock.Lock()
		for _, shardIdLocations := range resp.ShardIdLocations {
			shardId := erasure_coding.ShardId(shardIdLocations.ShardId)
			delete(ecVolume.ShardLocations, shardId)
			for _, loc := range shardIdLocations.Locations {
				ecVolume.ShardLocations[shardId] = append(ecVolume.ShardLocations[shardId], loc.Url)
			}
		}
		ecVolume.ShardLocationsRefreshTime = time.Now()
		ecVolume.ShardLocationsLock.Unlock()

		return nil
	})
	return
}

func (s *Store) readRemoteEcShardInterval(ctx context.Context, sourceDataNodes []string, vid needle.VolumeId, shardId erasure_coding.ShardId, buf []byte, offset int64) (n int, err error) {

	if len(sourceDataNodes) == 0 {
		return 0, fmt.Errorf("failed to find ec shard %d.%d", vid, shardId)
	}

	for _, sourceDataNode := range sourceDataNodes {
		glog.V(4).Infof("read remote ec shard %d.%d from %s", vid, shardId, sourceDataNode)
		n, err = s.doReadRemoteEcShardInterval(ctx, sourceDataNode, vid, shardId, buf, offset)
		if err == nil {
			return
		}
		glog.V(1).Infof("read remote ec shard %d.%d from %s: %v", vid, shardId, sourceDataNode, err)
	}

	return
}

func (s *Store) doReadRemoteEcShardInterval(ctx context.Context, sourceDataNode string, vid needle.VolumeId, shardId erasure_coding.ShardId, buf []byte, offset int64) (n int, err error) {

	err = operation.WithVolumeServerClient(sourceDataNode, s.grpcDialOption, func(client volume_server_pb.VolumeServerClient) error {

		// copy data slice
		shardReadClient, err := client.VolumeEcShardRead(ctx, &volume_server_pb.VolumeEcShardReadRequest{
			VolumeId: uint32(vid),
			ShardId:  uint32(shardId),
			Offset:   offset,
			Size:     int64(len(buf)),
		})
		if err != nil {
			return fmt.Errorf("failed to start reading ec shard %d.%d from %s: %v", vid, shardId, sourceDataNode, err)
		}

		for {
			resp, receiveErr := shardReadClient.Recv()
			if receiveErr == io.EOF {
				break
			}
			if receiveErr != nil {
				return fmt.Errorf("receiving ec shard %d.%d from %s: %v", vid, shardId, sourceDataNode, err)
			}
			copy(buf[n:n+len(resp.Data)], resp.Data)
			n += len(resp.Data)
		}

		return nil
	})
	if err != nil {
		return 0, fmt.Errorf("read ec shard %d.%d from %s: %v", vid, shardId, sourceDataNode, err)
	}

	return
}

func (s *Store) recoverOneRemoteEcShardInterval(ctx context.Context, ecVolume *erasure_coding.EcVolume, shardIdToRecover erasure_coding.ShardId, buf []byte, offset int64) (n int, err error) {
	glog.V(4).Infof("recover ec shard %d.%d from other locations", ecVolume.VolumeId, shardIdToRecover)

	enc, err := reedsolomon.New(erasure_coding.DataShardsCount, erasure_coding.ParityShardsCount)
	if err != nil {
		return 0, fmt.Errorf("failed to create encoder: %v", err)
	}

	bufs := make([][]byte, erasure_coding.TotalShardsCount)

	var wg sync.WaitGroup
	ecVolume.ShardLocationsLock.RLock()
	for shardId, locations := range ecVolume.ShardLocations {

		// skip currnent shard or empty shard
		if shardId == shardIdToRecover {
			continue
		}
		if len(locations) == 0 {
			glog.V(3).Infof("readRemoteEcShardInterval missing %d.%d from %+v", ecVolume.VolumeId, shardId, locations)
			continue
		}

		// read from remote locations
		wg.Add(1)
		go func(shardId erasure_coding.ShardId, locations []string) {
			defer wg.Done()
			data := make([]byte, len(buf))
			nRead, readErr := s.readRemoteEcShardInterval(ctx, locations, ecVolume.VolumeId, shardId, data, offset)
			if readErr != nil {
				glog.V(3).Infof("recover: readRemoteEcShardInterval %d.%d %d bytes from %+v: %v", ecVolume.VolumeId, shardId, nRead, locations, readErr)
				forgetShardId(ecVolume, shardId)
			}
			if nRead == len(buf) {
				bufs[shardId] = data
			}
		}(shardId, locations)
	}
	ecVolume.ShardLocationsLock.RUnlock()

	wg.Wait()

	if err = enc.ReconstructData(bufs); err != nil {
		glog.V(3).Infof("recovered ec shard %d.%d failed: %v", ecVolume.VolumeId, shardIdToRecover, err)
		return 0, err
	}
	glog.V(4).Infof("recovered ec shard %d.%d from other locations", ecVolume.VolumeId, shardIdToRecover)

	copy(buf, bufs[shardIdToRecover])

	return len(buf), nil
}
<|endoftext|>


<commit_before>package sgload

import (
	"log"
	"sync"
)

type Writer struct {
	UserCred
	ID                  int       // The numeric ID of the writer (ephemeral, only stored in memory)
	CreateDataStoreUser bool      // Whether this writer must first create a user on the DataStore service, ot just assume it already exists
	DataStore           DataStore // The target data store where docs will be written
	OutboundDocs        chan []Document
	WaitGroup           *sync.WaitGroup
	BatchSize           int
}

func NewWriter(wg *sync.WaitGroup, ID int, u UserCred, d DataStore, batchsize int) *Writer {

	outboundDocs := make(chan []Document, 100)

	return &Writer{
		UserCred:     u,
		ID:           ID,
		DataStore:    d,
		OutboundDocs: outboundDocs,
		WaitGroup:    wg,
		BatchSize:    batchsize,
	}
}

func (w *Writer) Run() {

	defer w.WaitGroup.Done()

	if w.CreateDataStoreUser == true {
		if err := w.DataStore.CreateUser(w.UserCred); err != nil {
			log.Fatalf("Error creating user in datastore.  User: %v, Err: %v", w.UserCred, err)
		}
	}

	for {

		select {
		case docs := <-w.OutboundDocs:

			switch len(docs) {
			case 1:
				doc := docs[0]
				_, ok := doc["_terminal"]
				if ok {
					return
				}

				if err := w.DataStore.CreateDocument(doc); err != nil {
					log.Fatalf("Error creating doc in datastore.  Doc: %v, Err: %v", doc, err)
				}

			default:
				if err := w.DataStore.BulkCreateDocuments(docs); err != nil {
					log.Fatalf("Error creating docs in datastore.  Docs: %v, Err: %v", docs, err)
				}

			}
		}

	}

}

func (w *Writer) AddToDataStore(docs []Document) {

	switch w.BatchSize {
	case 1:
		for _, doc := range docs {
			log.Printf("Writing doc to writer: %v", w)
			w.OutboundDocs <- []Document{doc}
			log.Printf("/Writing doc to writer: %v", w)
		}

	default:
		docBatches := breakIntoBatches(w.BatchSize, docs)
		for _, docBatch := range docBatches {
			w.OutboundDocs <- docBatch
		}
	}

}

// Break things into batches, for example:
//
// batchSize: 3
// things: [t1, t2, t3, t4, t5]
//
// result:
//
//   [
//     [t1, t2, t3],  <-- batch 1
//     [t4, t5],      <-- batch 2 (incomplete, not enough to fill batch)
//
//   ]
func breakIntoBatches(batchSize int, docs []Document) [][]Document {

	batches := [][]Document{}

	numBatches := len(docs) / batchSize

	// is there residue?  if so, add one more to batch
	if len(docs)%batchSize != 0 {
		numBatches += 1
	}

	for i := 0; i < numBatches; i++ {
		batch := []Document{}
		for j := 0; j < batchSize; j++ {
			docIndex := i*batchSize + j
			if docIndex >= len(docs) {
				break
			}
			doc := docs[docIndex]
			batch = append(batch, doc)
		}
		batches = append(batches, batch)
	}

	return batches

}
<commit_msg>Remove unwanted logging<commit_after>package sgload

import (
	"log"
	"sync"
)

type Writer struct {
	UserCred
	ID                  int       // The numeric ID of the writer (ephemeral, only stored in memory)
	CreateDataStoreUser bool      // Whether this writer must first create a user on the DataStore service, ot just assume it already exists
	DataStore           DataStore // The target data store where docs will be written
	OutboundDocs        chan []Document
	WaitGroup           *sync.WaitGroup
	BatchSize           int
}

func NewWriter(wg *sync.WaitGroup, ID int, u UserCred, d DataStore, batchsize int) *Writer {

	outboundDocs := make(chan []Document, 100)

	return &Writer{
		UserCred:     u,
		ID:           ID,
		DataStore:    d,
		OutboundDocs: outboundDocs,
		WaitGroup:    wg,
		BatchSize:    batchsize,
	}
}

func (w *Writer) Run() {

	defer w.WaitGroup.Done()

	if w.CreateDataStoreUser == true {
		if err := w.DataStore.CreateUser(w.UserCred); err != nil {
			log.Fatalf("Error creating user in datastore.  User: %v, Err: %v", w.UserCred, err)
		}
	}

	for {

		select {
		case docs := <-w.OutboundDocs:

			switch len(docs) {
			case 1:
				doc := docs[0]
				_, ok := doc["_terminal"]
				if ok {
					return
				}

				if err := w.DataStore.CreateDocument(doc); err != nil {
					log.Fatalf("Error creating doc in datastore.  Doc: %v, Err: %v", doc, err)
				}

			default:
				if err := w.DataStore.BulkCreateDocuments(docs); err != nil {
					log.Fatalf("Error creating docs in datastore.  Docs: %v, Err: %v", docs, err)
				}

			}
		}

	}

}

func (w *Writer) AddToDataStore(docs []Document) {

	switch w.BatchSize {
	case 1:
		for _, doc := range docs {
			w.OutboundDocs <- []Document{doc}
		}

	default:
		docBatches := breakIntoBatches(w.BatchSize, docs)
		for _, docBatch := range docBatches {
			w.OutboundDocs <- docBatch
		}
	}

}

// Break things into batches, for example:
//
// batchSize: 3
// things: [t1, t2, t3, t4, t5]
//
// result:
//
//   [
//     [t1, t2, t3],  <-- batch 1
//     [t4, t5],      <-- batch 2 (incomplete, not enough to fill batch)
//
//   ]
func breakIntoBatches(batchSize int, docs []Document) [][]Document {

	batches := [][]Document{}

	numBatches := len(docs) / batchSize

	// is there residue?  if so, add one more to batch
	if len(docs)%batchSize != 0 {
		numBatches += 1
	}

	for i := 0; i < numBatches; i++ {
		batch := []Document{}
		for j := 0; j < batchSize; j++ {
			docIndex := i*batchSize + j
			if docIndex >= len(docs) {
				break
			}
			doc := docs[docIndex]
			batch = append(batch, doc)
		}
		batches = append(batches, batch)
	}

	return batches

}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"log"
	"net/url"
	"os"
	"sort"
	"sync"

	"github.com/flynn/flynn-controller/client"
	ct "github.com/flynn/flynn-controller/types"
	"github.com/flynn/flynn-host/types"
	"github.com/flynn/go-dockerclient"
	"github.com/flynn/go-flynn/cluster"
)

func main() {
	cc, err := controller.New(os.Args[1])
	if err != nil {
		log.Fatal(err)
	}
	cl, err := cluster.NewClient()
	if err != nil {
		log.Fatal(err)
	}
	c := newContext(cl)
	// TODO: initial load of data
	// TODO: periodic full cluster sync for anti-entropy
	c.watchFormations(cc)
}

func newContext(cl clusterClient) *context {
	return &context{
		clusterClient: cl,
		formations:    NewFormations(),
		hosts:         newHostClients(),
		jobs:          newJobMap(),
	}
}

type context struct {
	clusterClient
	formations *Formations

	hosts *hostClients
	jobs  *jobMap
}

type clusterClient interface {
	ListHosts() (map[string]host.Host, error)
	AddJobs(req *host.AddJobsReq) (*host.AddJobsRes, error)
	ConnectHost(id string) (cluster.Host, error)
}

type formationStreamer interface {
	StreamFormations() (<-chan *ct.ExpandedFormation, *error)
}

func (c *context) watchFormations(fs formationStreamer) {
	ch, _ := fs.StreamFormations()

	for ef := range ch {
		f := NewFormation(c, ef)
		c.formations.Add(f)
		go f.Rectify()
	}

	// TODO: log disconnect and restart
	// TODO: trigger cluster sync
}

func (c *context) watchHost(id string) {
	if !c.hosts.Add(id) {
		return
	}
	defer c.hosts.Remove(id)

	h, err := c.ConnectHost(id)
	if err != nil {
		// TODO: log/handle error
	}
	c.hosts.Set(id, h)

	ch := make(chan *host.Event)
	h.StreamEvents("all", ch)
	for event := range ch {
		if event.Event != "error" && event.Event != "stop" {
			continue
		}
		job := c.jobs.Get(id, event.JobID)
		if job == nil {
			continue
		}

		c.jobs.Remove(id, event.JobID)
		go job.Formation.RemoveJob(job.Type, id, event.JobID)
	}
	// TODO: check error/reconnect
}

func newHostClients() *hostClients {
	return &hostClients{hosts: make(map[string]cluster.Host)}
}

type hostClients struct {
	hosts map[string]cluster.Host
	mtx   sync.RWMutex
}

func (h *hostClients) Add(id string) bool {
	h.mtx.Lock()
	defer h.mtx.Unlock()
	if _, exists := h.hosts[id]; exists {
		return false
	}
	h.hosts[id] = nil
	return true
}

func (h *hostClients) Set(id string, client cluster.Host) {
	h.mtx.Lock()
	h.hosts[id] = client
	h.mtx.Unlock()
}

func (h *hostClients) Remove(id string) {
	h.mtx.Lock()
	delete(h.hosts, id)
	h.mtx.Unlock()
}

func (h *hostClients) Get(id string) cluster.Host {
	h.mtx.RLock()
	defer h.mtx.RUnlock()
	return h.hosts[id]
}

func newJobMap() *jobMap {
	return &jobMap{jobs: make(map[jobKey]*Job)}
}

type jobMap struct {
	jobs map[jobKey]*Job
	mtx  sync.RWMutex
}

func (m *jobMap) Add(hostID, jobID string, job *Job) {
	m.mtx.Lock()
	m.jobs[jobKey{hostID, jobID}] = job
	m.mtx.Unlock()
}

func (m *jobMap) Remove(host, job string) {
	m.mtx.Lock()
	delete(m.jobs, jobKey{host, job})
	m.mtx.Unlock()
}

func (m *jobMap) Get(host, job string) *Job {
	m.mtx.RLock()
	defer m.mtx.RUnlock()
	return m.jobs[jobKey{host, job}]
}

type jobKey struct {
	hostID, jobID string
}

type formationKey struct {
	appID, releaseID string
}

func NewFormations() *Formations {
	return &Formations{formations: make(map[formationKey]*Formation)}
}

type Formations struct {
	formations map[formationKey]*Formation
	mtx        sync.RWMutex
}

func (fs *Formations) Get(appID, releaseID string) *Formation {
	fs.mtx.RLock()
	defer fs.mtx.RUnlock()
	return fs.formations[formationKey{appID, releaseID}]
}

func (fs *Formations) Add(f *Formation) {
	fs.mtx.Lock()
	fs.formations[f.key()] = f
	fs.mtx.Unlock()
}

func (fs *Formations) Delete(f *Formation) {
	fs.mtx.Lock()
	delete(fs.formations, f.key())
	fs.mtx.Unlock()
}

func NewFormation(c *context, ef *ct.ExpandedFormation) *Formation {
	return &Formation{
		App:       ef.App,
		Release:   ef.Release,
		Artifact:  ef.Artifact,
		Processes: ef.Processes,
		jobs:      make(jobTypeMap),
		c:         c,
	}
}

type Job struct {
	Type      string
	Formation *Formation
}

type jobTypeMap map[string]map[jobKey]*Job

func (m jobTypeMap) Add(typ, host, id string) *Job {
	jobs, ok := m[typ]
	if !ok {
		jobs = make(map[jobKey]*Job)
		m[typ] = jobs
	}
	job := &Job{Type: typ}
	jobs[jobKey{host, id}] = job
	return job
}

func (m jobTypeMap) Remove(typ, host, id string) {
	if jobs, ok := m[typ]; ok {
		delete(jobs, jobKey{host, id})
	}
}

func (m jobTypeMap) Get(typ, host, id string) *Job {
	return m[typ][jobKey{host, id}]
}

type Formation struct {
	mtx       sync.Mutex
	App       *ct.App
	Release   *ct.Release
	Artifact  *ct.Artifact
	Processes map[string]int

	jobs jobTypeMap
	c    *context
}

func (f *Formation) key() formationKey {
	return formationKey{f.App.ID, f.Release.ID}
}

func (f *Formation) Rectify() {
	f.mtx.Lock()
	defer f.mtx.Unlock()
	f.rectify()
}

func (f *Formation) RemoveJob(typ, hostID, jobID string) {
	f.mtx.Lock()
	defer f.mtx.Unlock()

	f.jobs.Remove(typ, hostID, jobID)
	f.rectify()
}

func (f *Formation) rectify() {
	// update job counts
	for t, expected := range f.Processes {
		diff := expected - len(f.jobs[t])
		if diff > 0 {
			f.add(diff, t)
		} else if diff < 0 {
			f.remove(-diff, t)
		}
	}

	// remove process types
	for t, jobs := range f.jobs {
		if _, exists := f.Processes[t]; !exists {
			f.remove(len(jobs), t)
		}
	}
}

func (f *Formation) add(n int, name string) {
	config, err := f.jobConfig(name)
	if err != nil {
		// TODO: log/handle error
	}
	for i := 0; i < n; i++ {
		hosts, err := f.c.ListHosts()
		if err != nil {
			// TODO: log/handle error
		}
		if len(hosts) == 0 {
			// TODO: log/handle error
		}
		hostCounts := make(map[string]int, len(hosts))
		for _, h := range hosts {
			hostCounts[h.ID] = 0
			for _, job := range h.Jobs {
				if f.jobType(job) != name {
					continue
				}
				hostCounts[h.ID]++
			}
		}
		sh := make(sortHosts, 0, len(hosts))
		for id, count := range hostCounts {
			sh = append(sh, sortHost{id, count})
		}
		sh.Sort()

		h := hosts[sh[0].ID]
		go f.c.watchHost(h.ID)

		job := f.jobs.Add(name, h.ID, config.ID)
		job.Formation = f
		f.c.jobs.Add(h.ID, config.ID, job)

		res, err := f.c.AddJobs(&host.AddJobsReq{HostJobs: map[string][]*host.Job{h.ID: {config}}})
		if err != nil || !res.Success {
			f.jobs.Remove(name, h.ID, config.ID)
			f.c.jobs.Remove(h.ID, config.ID)
			// TODO: log/handle error
		}
	}
}

func (f *Formation) jobType(job *host.Job) string {
	if job.Attributes["flynn-controller.app"] != f.App.ID ||
		job.Attributes["flynn-controller.release"] != f.Release.ID {
		return ""
	}
	return job.Attributes["flynn-controller.type"]
}

func (f *Formation) remove(n int, name string) {
	i := 0
	for k := range f.jobs[name] {
		// TODO: robust host handling
		if err := f.c.hosts.Get(k.hostID).StopJob(k.jobID); err != nil {
			// TODO: log/handle error
		}
		f.jobs.Remove(name, k.hostID, k.jobID)
		f.c.jobs.Remove(k.hostID, k.jobID)
		if i++; i == n {
			break
		}
	}
}

func (f *Formation) jobConfig(name string) (*host.Job, error) {
	t := f.Release.Processes[name]
	image, err := dockerImage(f.Artifact.URI)
	if err != nil {
		return nil, err
	}
	return &host.Job{
		ID:       cluster.RandomJobID(""),
		TCPPorts: t.Ports.TCP,
		Attributes: map[string]string{
			"flynn-controller.app":     f.App.ID,
			"flynn-controller.release": f.Release.ID,
			"flynn-controller.type":    name,
		},
		Config: &docker.Config{
			Cmd:   t.Cmd,
			Env:   formatEnv(f.Release.Env, t.Env),
			Image: image,
		},
	}, nil
}

func dockerImage(uri string) (string, error) {
	// TODO: ID refs (see https://github.com/dotcloud/docker/issues/4106)
	u, err := url.Parse(uri)
	if err != nil {
		return "", err
	}
	if u.Scheme != "docker" {
		return "", errors.New("scheduler: only docker artifact URIs are currently supported")
	}
	var suffix string
	if tag := u.Query().Get("tag"); tag != "" {
		suffix = ":" + tag
	}
	return u.Host + u.Path + suffix, nil
}

func formatEnv(envs ...map[string]string) []string {
	env := make(map[string]string)
	for _, e := range envs {
		for k, v := range e {
			env[k] = v
		}
	}
	res := make([]string, 0, len(env))
	for k, v := range env {
		res = append(res, k+"="+v)
	}
	return res
}

type sortHost struct {
	ID   string
	Jobs int
}

type sortHosts []sortHost

func (h sortHosts) Len() int           { return len(h) }
func (h sortHosts) Less(i, j int) bool { return h[i].Jobs < h[j].Jobs }
func (h sortHosts) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h sortHosts) Sort()              { sort.Sort(h) }
<commit_msg>controller/scheduler: Fix formation updates<commit_after>package main

import (
	"errors"
	"log"
	"net/url"
	"os"
	"sort"
	"sync"

	"github.com/flynn/flynn-controller/client"
	ct "github.com/flynn/flynn-controller/types"
	"github.com/flynn/flynn-host/types"
	"github.com/flynn/go-dockerclient"
	"github.com/flynn/go-flynn/cluster"
)

func main() {
	cc, err := controller.New(os.Args[1])
	if err != nil {
		log.Fatal(err)
	}
	cl, err := cluster.NewClient()
	if err != nil {
		log.Fatal(err)
	}
	c := newContext(cl)
	// TODO: initial load of data
	// TODO: periodic full cluster sync for anti-entropy
	c.watchFormations(cc)
}

func newContext(cl clusterClient) *context {
	return &context{
		clusterClient: cl,
		formations:    NewFormations(),
		hosts:         newHostClients(),
		jobs:          newJobMap(),
	}
}

type context struct {
	clusterClient
	formations *Formations

	hosts *hostClients
	jobs  *jobMap
}

type clusterClient interface {
	ListHosts() (map[string]host.Host, error)
	AddJobs(req *host.AddJobsReq) (*host.AddJobsRes, error)
	ConnectHost(id string) (cluster.Host, error)
}

type formationStreamer interface {
	StreamFormations() (<-chan *ct.ExpandedFormation, *error)
}

func (c *context) watchFormations(fs formationStreamer) {
	ch, _ := fs.StreamFormations()

	for ef := range ch {
		f := c.formations.Get(ef.App.ID, ef.Release.ID)
		if f != nil {
			f.SetProcesses(ef.Processes)
		} else {
			f = NewFormation(c, ef)
			c.formations.Add(f)
		}
		go f.Rectify()
	}

	// TODO: log disconnect and restart
	// TODO: trigger cluster sync
}

func (c *context) watchHost(id string) {
	if !c.hosts.Add(id) {
		return
	}
	defer c.hosts.Remove(id)

	h, err := c.ConnectHost(id)
	if err != nil {
		// TODO: log/handle error
	}
	c.hosts.Set(id, h)

	ch := make(chan *host.Event)
	h.StreamEvents("all", ch)
	for event := range ch {
		if event.Event != "error" && event.Event != "stop" {
			continue
		}
		job := c.jobs.Get(id, event.JobID)
		if job == nil {
			continue
		}

		c.jobs.Remove(id, event.JobID)
		go job.Formation.RemoveJob(job.Type, id, event.JobID)
	}
	// TODO: check error/reconnect
}

func newHostClients() *hostClients {
	return &hostClients{hosts: make(map[string]cluster.Host)}
}

type hostClients struct {
	hosts map[string]cluster.Host
	mtx   sync.RWMutex
}

func (h *hostClients) Add(id string) bool {
	h.mtx.Lock()
	defer h.mtx.Unlock()
	if _, exists := h.hosts[id]; exists {
		return false
	}
	h.hosts[id] = nil
	return true
}

func (h *hostClients) Set(id string, client cluster.Host) {
	h.mtx.Lock()
	h.hosts[id] = client
	h.mtx.Unlock()
}

func (h *hostClients) Remove(id string) {
	h.mtx.Lock()
	delete(h.hosts, id)
	h.mtx.Unlock()
}

func (h *hostClients) Get(id string) cluster.Host {
	h.mtx.RLock()
	defer h.mtx.RUnlock()
	return h.hosts[id]
}

func newJobMap() *jobMap {
	return &jobMap{jobs: make(map[jobKey]*Job)}
}

type jobMap struct {
	jobs map[jobKey]*Job
	mtx  sync.RWMutex
}

func (m *jobMap) Add(hostID, jobID string, job *Job) {
	m.mtx.Lock()
	m.jobs[jobKey{hostID, jobID}] = job
	m.mtx.Unlock()
}

func (m *jobMap) Remove(host, job string) {
	m.mtx.Lock()
	delete(m.jobs, jobKey{host, job})
	m.mtx.Unlock()
}

func (m *jobMap) Get(host, job string) *Job {
	m.mtx.RLock()
	defer m.mtx.RUnlock()
	return m.jobs[jobKey{host, job}]
}

type jobKey struct {
	hostID, jobID string
}

type formationKey struct {
	appID, releaseID string
}

func NewFormations() *Formations {
	return &Formations{formations: make(map[formationKey]*Formation)}
}

type Formations struct {
	formations map[formationKey]*Formation
	mtx        sync.RWMutex
}

func (fs *Formations) Get(appID, releaseID string) *Formation {
	fs.mtx.RLock()
	defer fs.mtx.RUnlock()
	return fs.formations[formationKey{appID, releaseID}]
}

func (fs *Formations) Add(f *Formation) {
	fs.mtx.Lock()
	fs.formations[f.key()] = f
	fs.mtx.Unlock()
}

func (fs *Formations) Delete(f *Formation) {
	fs.mtx.Lock()
	delete(fs.formations, f.key())
	fs.mtx.Unlock()
}

func NewFormation(c *context, ef *ct.ExpandedFormation) *Formation {
	return &Formation{
		App:       ef.App,
		Release:   ef.Release,
		Artifact:  ef.Artifact,
		Processes: ef.Processes,
		jobs:      make(jobTypeMap),
		c:         c,
	}
}

type Job struct {
	Type      string
	Formation *Formation
}

type jobTypeMap map[string]map[jobKey]*Job

func (m jobTypeMap) Add(typ, host, id string) *Job {
	jobs, ok := m[typ]
	if !ok {
		jobs = make(map[jobKey]*Job)
		m[typ] = jobs
	}
	job := &Job{Type: typ}
	jobs[jobKey{host, id}] = job
	return job
}

func (m jobTypeMap) Remove(typ, host, id string) {
	if jobs, ok := m[typ]; ok {
		delete(jobs, jobKey{host, id})
	}
}

func (m jobTypeMap) Get(typ, host, id string) *Job {
	return m[typ][jobKey{host, id}]
}

type Formation struct {
	mtx       sync.Mutex
	App       *ct.App
	Release   *ct.Release
	Artifact  *ct.Artifact
	Processes map[string]int

	jobs jobTypeMap
	c    *context
}

func (f *Formation) key() formationKey {
	return formationKey{f.App.ID, f.Release.ID}
}

func (f *Formation) SetProcesses(p map[string]int) {
	f.mtx.Lock()
	f.Processes = p
	f.mtx.Unlock()
}

func (f *Formation) Rectify() {
	f.mtx.Lock()
	defer f.mtx.Unlock()
	f.rectify()
}

func (f *Formation) RemoveJob(typ, hostID, jobID string) {
	f.mtx.Lock()
	defer f.mtx.Unlock()

	f.jobs.Remove(typ, hostID, jobID)
	f.rectify()
}

func (f *Formation) rectify() {
	// update job counts
	for t, expected := range f.Processes {
		diff := expected - len(f.jobs[t])
		if diff > 0 {
			f.add(diff, t)
		} else if diff < 0 {
			f.remove(-diff, t)
		}
	}

	// remove process types
	for t, jobs := range f.jobs {
		if _, exists := f.Processes[t]; !exists {
			f.remove(len(jobs), t)
		}
	}
}

func (f *Formation) add(n int, name string) {
	config, err := f.jobConfig(name)
	if err != nil {
		// TODO: log/handle error
	}
	for i := 0; i < n; i++ {
		hosts, err := f.c.ListHosts()
		if err != nil {
			// TODO: log/handle error
		}
		if len(hosts) == 0 {
			// TODO: log/handle error
		}
		hostCounts := make(map[string]int, len(hosts))
		for _, h := range hosts {
			hostCounts[h.ID] = 0
			for _, job := range h.Jobs {
				if f.jobType(job) != name {
					continue
				}
				hostCounts[h.ID]++
			}
		}
		sh := make(sortHosts, 0, len(hosts))
		for id, count := range hostCounts {
			sh = append(sh, sortHost{id, count})
		}
		sh.Sort()

		h := hosts[sh[0].ID]
		go f.c.watchHost(h.ID)

		job := f.jobs.Add(name, h.ID, config.ID)
		job.Formation = f
		f.c.jobs.Add(h.ID, config.ID, job)

		res, err := f.c.AddJobs(&host.AddJobsReq{HostJobs: map[string][]*host.Job{h.ID: {config}}})
		if err != nil || !res.Success {
			f.jobs.Remove(name, h.ID, config.ID)
			f.c.jobs.Remove(h.ID, config.ID)
			// TODO: log/handle error
		}
	}
}

func (f *Formation) jobType(job *host.Job) string {
	if job.Attributes["flynn-controller.app"] != f.App.ID ||
		job.Attributes["flynn-controller.release"] != f.Release.ID {
		return ""
	}
	return job.Attributes["flynn-controller.type"]
}

func (f *Formation) remove(n int, name string) {
	i := 0
	for k := range f.jobs[name] {
		// TODO: robust host handling
		if err := f.c.hosts.Get(k.hostID).StopJob(k.jobID); err != nil {
			// TODO: log/handle error
		}
		f.jobs.Remove(name, k.hostID, k.jobID)
		f.c.jobs.Remove(k.hostID, k.jobID)
		if i++; i == n {
			break
		}
	}
}

func (f *Formation) jobConfig(name string) (*host.Job, error) {
	t := f.Release.Processes[name]
	image, err := dockerImage(f.Artifact.URI)
	if err != nil {
		return nil, err
	}
	return &host.Job{
		ID:       cluster.RandomJobID(""),
		TCPPorts: t.Ports.TCP,
		Attributes: map[string]string{
			"flynn-controller.app":     f.App.ID,
			"flynn-controller.release": f.Release.ID,
			"flynn-controller.type":    name,
		},
		Config: &docker.Config{
			Cmd:   t.Cmd,
			Env:   formatEnv(f.Release.Env, t.Env),
			Image: image,
		},
	}, nil
}

func dockerImage(uri string) (string, error) {
	// TODO: ID refs (see https://github.com/dotcloud/docker/issues/4106)
	u, err := url.Parse(uri)
	if err != nil {
		return "", err
	}
	if u.Scheme != "docker" {
		return "", errors.New("scheduler: only docker artifact URIs are currently supported")
	}
	var suffix string
	if tag := u.Query().Get("tag"); tag != "" {
		suffix = ":" + tag
	}
	return u.Host + u.Path + suffix, nil
}

func formatEnv(envs ...map[string]string) []string {
	env := make(map[string]string)
	for _, e := range envs {
		for k, v := range e {
			env[k] = v
		}
	}
	res := make([]string, 0, len(env))
	for k, v := range env {
		res = append(res, k+"="+v)
	}
	return res
}

type sortHost struct {
	ID   string
	Jobs int
}

type sortHosts []sortHost

func (h sortHosts) Len() int           { return len(h) }
func (h sortHosts) Less(i, j int) bool { return h[i].Jobs < h[j].Jobs }
func (h sortHosts) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h sortHosts) Sort()              { sort.Sort(h) }
<|endoftext|>


<commit_before>package config

import (
	"encoding/json"
	"fmt"
	"os"
	"reflect"
	"strings"
)

// ErrPrimaryConfigFileNotExist is returned if the primary JSON config file doesn't exist
var ErrPrimaryConfigFileNotExist = fmt.Errorf("config: primary config file does not exist")

// Load will load a configuration json file into a struct
func Load(path, environment string, configData interface{}) (err error) {
	file, err := os.Open(path)
	if os.IsNotExist(err) {
		return ErrPrimaryConfigFileNotExist
	} else if err != nil {
		return err
	}

	err = json.NewDecoder(file).Decode(configData)
	if err != nil {
		return err
	}

	altPath := strings.Replace(path, ".json", "."+environment+".json", 1)

	altFile, err := os.Open(altPath)
	if err != nil && !os.IsNotExist(err) {
		return err
	}

	if os.IsNotExist(err) {
		return nil
	}

	altConfigData := map[string]interface{}{}
	err = json.NewDecoder(altFile).Decode(&altConfigData)
	if err != nil {
		return err
	}

	defer func() {
		if p := recover(); p != nil {
			err = fmt.Errorf("config: parsing environment config file: %s", p)
		}
	}()

	configValue := reflect.ValueOf(configData).Elem()
	parseMap(altConfigData, configValue)

	return
}

func parseMap(aMap map[string]interface{}, configValue reflect.Value) {
	for key, value := range aMap {
		fieldName := ""

		for i := 0; i < configValue.NumField(); i++ {
			fieldInfo := configValue.Type().Field(i)
			jsonFieldName := strings.TrimSpace(fieldInfo.Tag.Get("json"))

			if jsonFieldName == key {
				fieldName = fieldInfo.Name
			}
		}

		fieldValue := configValue.FieldByName(fieldName)

		switch realValue := value.(type) {
		case map[string]interface{}:
			if fieldValue.Kind() == reflect.Struct || fieldValue.Kind() == reflect.Map {
				parseMap(realValue, fieldValue)
			}
		case []interface{}:
			if fieldValue.Kind() == reflect.Slice {
				parseSlice(realValue, fieldValue)
			}
		case string:
			if fieldValue.Kind() == reflect.String {
				fieldValue.SetString(realValue)
			}
		case float64:
			switch fieldValue.Kind() {
			case reflect.Float32, reflect.Float64:
				fieldValue.SetFloat(realValue)
			case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
				fieldValue.SetInt(int64(realValue))
			case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
				fieldValue.SetUint(uint64(realValue))
			}
		case bool:
			if fieldValue.Kind() == reflect.Bool {
				fieldValue.SetBool(realValue)
			}
		}
	}
}

func parseSlice(aSlice []interface{}, configValue reflect.Value) {
	newSlice := reflect.MakeSlice(configValue.Type(), len(aSlice), len(aSlice))
	configValue.Set(newSlice)

	for i, value := range aSlice {
		switch realItem := value.(type) {
		case map[string]interface{}:
			parseMap(realItem, configValue.Index(i))
		case []interface{}:
			parseSlice(realItem, configValue.Index(i))
		case string:
			configValue.Index(i).SetString(realItem)
		case float64:
			switch configValue.Index(i).Kind() {
			case reflect.Float32, reflect.Float64:
				configValue.Index(i).SetFloat(realItem)
			case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
				configValue.Index(i).SetInt(int64(realItem))
			case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
				configValue.Index(i).SetUint(uint64(realItem))
			}
		case bool:
			configValue.Index(i).SetBool(realItem)
		}
	}
}
<commit_msg>added: more robust type checking and error handling for config lib.<commit_after>// Package config provides utilities for loading a JSON configuration file into a struct object
// graph, with support for providing an alternative environment JSON config file
// (e.g. "dev", "staging", "uat", "live"), with values replaced using transformations. Inspired by
// the way Microsoft ASP.NET handles configuration files.
package config

import (
	"encoding/json"
	"fmt"
	"os"
	"reflect"
	"strings"
)

// ErrPrimaryConfigFileNotExist is returned if the primary JSON config file doesn't exist
var ErrPrimaryConfigFileNotExist = fmt.Errorf("config: primary config file does not exist")

// ErrConfigDataNotPointer is returned when the configData
// struct to pass config data into is not a pointer
var ErrConfigDataNotPointer = fmt.Errorf("config: configData argument is not a pointer")

// Load will load a configuration json file into a struct
func Load(path, environment string, configData interface{}) (err error) {
	file, err := os.Open(path)
	if os.IsNotExist(err) {
		return ErrPrimaryConfigFileNotExist
	} else if err != nil {
		return fmt.Errorf("config: error opening primary config file: %s", err)
	}

	if reflect.TypeOf(configData).Kind() != reflect.Ptr {
		return ErrConfigDataNotPointer
	}

	err = json.NewDecoder(file).Decode(configData)
	if err != nil {
		return fmt.Errorf("config: cannot unmarshal config file: %s", err)
	}

	altPath := strings.Replace(path, ".json", "."+environment+".json", 1)

	altFile, err := os.Open(altPath)
	if err != nil && !os.IsNotExist(err) {
		return fmt.Errorf("config: error opening environment config file \"%s\": %s", altPath, err)
	}

	if os.IsNotExist(err) {
		return nil
	}

	altConfigData := map[string]interface{}{}
	err = json.NewDecoder(altFile).Decode(&altConfigData)
	if err != nil {
		return fmt.Errorf("config: cannot unmarshal environment config file: %s", err)
	}

	defer func() {
		if p := recover(); p != nil {
			err = fmt.Errorf("config: parsing environment config file: %s", p)
		}
	}()

	configValue := reflect.ValueOf(configData).Elem()
	parseMap(altConfigData, configValue)

	return
}

func parseMap(aMap map[string]interface{}, configValue reflect.Value) {
	for key, value := range aMap {
		fieldName := ""

		for i := 0; i < configValue.NumField(); i++ {
			fieldInfo := configValue.Type().Field(i)
			jsonFieldName := strings.TrimSpace(fieldInfo.Tag.Get("json"))

			if jsonFieldName == key {
				fieldName = fieldInfo.Name
			}
		}

		fieldValue := configValue.FieldByName(fieldName)
		if fieldValue.Kind() == reflect.Invalid {
			continue
		}

		switch realValue := value.(type) {
		case map[string]interface{}:
			if fieldValue.Kind() == reflect.Struct || fieldValue.Kind() == reflect.Map {
				parseMap(realValue, fieldValue)
			}
		case []interface{}:
			if fieldValue.Kind() == reflect.Slice {
				parseSlice(realValue, fieldValue)
			}
		case string:
			if fieldValue.Kind() == reflect.String {
				fieldValue.SetString(realValue)
			}
		case float64:
			switch fieldValue.Kind() {
			case reflect.Float32, reflect.Float64:
				fieldValue.SetFloat(realValue)
			case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
				fieldValue.SetInt(int64(realValue))
			case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
				fieldValue.SetUint(uint64(realValue))
			}
		case bool:
			if fieldValue.Kind() == reflect.Bool {
				fieldValue.SetBool(realValue)
			}
		}
	}
}

func parseSlice(aSlice []interface{}, configValue reflect.Value) {
	newSlice := reflect.MakeSlice(configValue.Type(), len(aSlice), len(aSlice))
	configValue.Set(newSlice)

	for i, value := range aSlice {
		configItem := configValue.Index(i)

		switch realItem := value.(type) {
		case map[string]interface{}:
			if configItem.Kind() == reflect.Struct || configItem.Kind() == reflect.Map {
				parseMap(realItem, configItem)
			}
		case []interface{}:
			if configItem.Kind() == reflect.Slice {
				parseSlice(realItem, configItem)
			}
		case string:
			if configItem.Kind() == reflect.String {
				configItem.SetString(realItem)
			}
		case float64:
			switch configItem.Kind() {
			case reflect.Float32, reflect.Float64:
				configItem.SetFloat(realItem)
			case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
				configItem.SetInt(int64(realItem))
			case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
				configItem.SetUint(uint64(realItem))
			}
		case bool:
			if configItem.Kind() == reflect.Bool {
				configItem.SetBool(realItem)
			}
		}
	}
}
<|endoftext|>


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

package audio

import (
	"fmt"
	"io"
	"sync"
	"time"
)

// TODO: The term 'buffer' is confusing. Name each buffer with good terms.

// oneBufferSize returns the size of one buffer in the player implementation.
func oneBufferSize(sampleRate int) int {
	return sampleRate * channelNum * bitDepthInBytes / 4
}

// maxBufferSize returns the maximum size of the buffer for the audio source.
// This buffer is used when unreading on pausing the player.
func maxBufferSize(sampleRate int) int {
	// Actually *2 should be enough in most cases,
	// but in some implementation (e.g, go2cpp), a player might have more UnplayedBufferSize values.
	// As a safe margin, use *4 value.
	// TODO: Ensure the maximum value of UnplayedBufferSize on all the platforms.
	return oneBufferSize(sampleRate) * 4
}

// readerDriver represents a driver using io.ReadClosers.
type readerDriver interface {
	NewPlayer(io.Reader) readerDriverPlayer
	io.Closer
}

type readerDriverPlayer interface {
	Pause()
	Play()
	IsPlaying() bool
	Reset()
	Volume() float64
	SetVolume(volume float64)
	UnplayedBufferSize() int64
	io.Closer
}

type readerPlayerFactory struct {
	driver     readerDriver
	sampleRate int
}

var readerDriverForTesting readerDriver

func newReaderPlayerFactory(sampleRate int) *readerPlayerFactory {
	f := &readerPlayerFactory{
		sampleRate: sampleRate,
	}
	if readerDriverForTesting != nil {
		f.driver = readerDriverForTesting
	}
	// TODO: Consider the hooks.
	return f
}

type readerPlayer struct {
	context *Context
	player  readerDriverPlayer
	stream  *timeStream
	factory *readerPlayerFactory
	m       sync.Mutex
}

func (f *readerPlayerFactory) newPlayerImpl(context *Context, src io.Reader) (playerImpl, error) {
	sampleRate := context.SampleRate()
	s, err := newTimeStream(src, sampleRate)
	if err != nil {
		return nil, err
	}

	p := &readerPlayer{
		context: context,
		stream:  s,
		factory: f,
	}
	return p, nil
}

func (p *readerPlayer) ensurePlayer() error {
	// Initialize the underlying player lazily to enable calling NewContext in an 'init' function.
	// Accessing the underlying player functions requires the environment to be already initialized,
	// but if Ebiten is used for a shared library, the timing when init functions are called
	// is unexpectable.
	// e.g. a variable for JVM on Android might not be set.
	if p.factory.driver == nil {
		d, err := newReaderDriverImpl(p.context)
		if err != nil {
			return err
		}
		p.factory.driver = d
	}
	if p.player == nil {
		p.player = p.factory.driver.NewPlayer(p.stream)
	}
	return nil
}

func (p *readerPlayer) Play() {
	p.m.Lock()
	defer p.m.Unlock()
	if err := p.ensurePlayer(); err != nil {
		p.context.setError(err)
		return
	}

	p.player.Play()
	p.context.addPlayer(p)
}

func (p *readerPlayer) Pause() {
	p.m.Lock()
	defer p.m.Unlock()
	if err := p.ensurePlayer(); err != nil {
		p.context.setError(err)
		return
	}

	n := p.player.UnplayedBufferSize()
	p.player.Pause()
	p.stream.Unread(int(n))
	p.context.removePlayer(p)
}

func (p *readerPlayer) IsPlaying() bool {
	p.m.Lock()
	defer p.m.Unlock()
	if err := p.ensurePlayer(); err != nil {
		p.context.setError(err)
		return false
	}

	return p.player.IsPlaying()
}

func (p *readerPlayer) Volume() float64 {
	p.m.Lock()
	defer p.m.Unlock()
	if err := p.ensurePlayer(); err != nil {
		p.context.setError(err)
		return 0
	}

	return p.player.Volume()
}

func (p *readerPlayer) SetVolume(volume float64) {
	p.m.Lock()
	defer p.m.Unlock()
	if err := p.ensurePlayer(); err != nil {
		p.context.setError(err)
		return
	}

	p.player.SetVolume(volume)
}

func (p *readerPlayer) Close() error {
	p.m.Lock()
	defer p.m.Unlock()

	if p.player != nil {
		p.player.Pause()
		return p.player.Close()
	}
	return nil
}

func (p *readerPlayer) Current() time.Duration {
	p.m.Lock()
	defer p.m.Unlock()
	if err := p.ensurePlayer(); err != nil {
		p.context.setError(err)
		return 0
	}

	sample := (p.stream.Current() - p.player.UnplayedBufferSize()) / bytesPerSample
	return time.Duration(sample) * time.Second / time.Duration(p.factory.sampleRate)
}

func (p *readerPlayer) Rewind() error {
	return p.Seek(0)
}

func (p *readerPlayer) Seek(offset time.Duration) error {
	p.m.Lock()
	defer p.m.Unlock()

	if p.player != nil {
		if p.player.IsPlaying() {
			defer func() {
				p.player.Play()
			}()
		}
		p.player.Reset()
	}
	return p.stream.Seek(offset)
}

func (p *readerPlayer) source() io.Reader {
	return p.stream.r
}

type timeStream struct {
	r          io.Reader
	sampleRate int
	pos        int64
	buf        []byte
	unread     int
}

func newTimeStream(r io.Reader, sampleRate int) (*timeStream, error) {
	s := &timeStream{
		r:          r,
		sampleRate: sampleRate,
	}
	if seeker, ok := s.r.(io.Seeker); ok {
		// Get the current position of the source.
		pos, err := seeker.Seek(0, io.SeekCurrent)
		if err != nil {
			return nil, err
		}
		s.pos = pos
	}
	return s, nil
}

func (s *timeStream) Unread(n int) {
	if s.unread+n > len(s.buf) {
		panic(fmt.Sprintf("audio: too much unreading: %d, the buffer size: %d, unreading position: %d", n, len(s.buf), s.unread))
	}
	s.unread += n
	s.pos -= int64(n)
}

func (s *timeStream) Read(buf []byte) (int, error) {
	if s.unread > 0 {
		n := copy(buf, s.buf[len(s.buf)-s.unread:])
		s.unread -= n
		s.pos += int64(n)
		return n, nil
	}

	n, err := s.r.Read(buf)
	s.pos += int64(n)
	s.buf = append(s.buf, buf[:n]...)
	if m := maxBufferSize(s.sampleRate); len(s.buf) > m {
		s.buf = s.buf[len(s.buf)-m:]
	}
	return n, err
}

func (s *timeStream) Seek(offset time.Duration) error {
	o := int64(offset) * bytesPerSample * int64(s.sampleRate) / int64(time.Second)

	// Align the byte position with the samples.
	o -= o % bytesPerSample
	o += s.pos % bytesPerSample

	seeker, ok := s.r.(io.Seeker)
	if !ok {
		panic("audio: the source must be io.Seeker when seeking but not")
	}
	pos, err := seeker.Seek(o, io.SeekStart)
	if err != nil {
		return err
	}

	s.pos = pos
	s.buf = s.buf[:0]
	s.unread = 0
	return nil
}

func (s *timeStream) Current() int64 {
	return s.pos
}
<commit_msg>audio: Ensure Close is called at a readerPlayer is GCed<commit_after>// Copyright 2021 The Ebiten Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package audio

import (
	"fmt"
	"io"
	"runtime"
	"sync"
	"time"
)

// TODO: The term 'buffer' is confusing. Name each buffer with good terms.

// oneBufferSize returns the size of one buffer in the player implementation.
func oneBufferSize(sampleRate int) int {
	return sampleRate * channelNum * bitDepthInBytes / 4
}

// maxBufferSize returns the maximum size of the buffer for the audio source.
// This buffer is used when unreading on pausing the player.
func maxBufferSize(sampleRate int) int {
	// Actually *2 should be enough in most cases,
	// but in some implementation (e.g, go2cpp), a player might have more UnplayedBufferSize values.
	// As a safe margin, use *4 value.
	// TODO: Ensure the maximum value of UnplayedBufferSize on all the platforms.
	return oneBufferSize(sampleRate) * 4
}

// readerDriver represents a driver using io.ReadClosers.
type readerDriver interface {
	NewPlayer(io.Reader) readerDriverPlayer
	io.Closer
}

type readerDriverPlayer interface {
	Pause()
	Play()
	IsPlaying() bool
	Reset()
	Volume() float64
	SetVolume(volume float64)
	UnplayedBufferSize() int64
	io.Closer
}

type readerPlayerFactory struct {
	driver     readerDriver
	sampleRate int
}

var readerDriverForTesting readerDriver

func newReaderPlayerFactory(sampleRate int) *readerPlayerFactory {
	f := &readerPlayerFactory{
		sampleRate: sampleRate,
	}
	if readerDriverForTesting != nil {
		f.driver = readerDriverForTesting
	}
	// TODO: Consider the hooks.
	return f
}

type readerPlayer struct {
	context *Context
	player  readerDriverPlayer
	stream  *timeStream
	factory *readerPlayerFactory
	m       sync.Mutex
}

func (f *readerPlayerFactory) newPlayerImpl(context *Context, src io.Reader) (playerImpl, error) {
	sampleRate := context.SampleRate()
	s, err := newTimeStream(src, sampleRate)
	if err != nil {
		return nil, err
	}

	p := &readerPlayer{
		context: context,
		stream:  s,
		factory: f,
	}
	runtime.SetFinalizer(p, (*readerPlayer).Close)
	return p, nil
}

func (p *readerPlayer) ensurePlayer() error {
	// Initialize the underlying player lazily to enable calling NewContext in an 'init' function.
	// Accessing the underlying player functions requires the environment to be already initialized,
	// but if Ebiten is used for a shared library, the timing when init functions are called
	// is unexpectable.
	// e.g. a variable for JVM on Android might not be set.
	if p.factory.driver == nil {
		d, err := newReaderDriverImpl(p.context)
		if err != nil {
			return err
		}
		p.factory.driver = d
	}
	if p.player == nil {
		p.player = p.factory.driver.NewPlayer(p.stream)
	}
	return nil
}

func (p *readerPlayer) Play() {
	p.m.Lock()
	defer p.m.Unlock()
	if err := p.ensurePlayer(); err != nil {
		p.context.setError(err)
		return
	}

	p.player.Play()
	p.context.addPlayer(p)
}

func (p *readerPlayer) Pause() {
	p.m.Lock()
	defer p.m.Unlock()
	if err := p.ensurePlayer(); err != nil {
		p.context.setError(err)
		return
	}

	n := p.player.UnplayedBufferSize()
	p.player.Pause()
	p.stream.Unread(int(n))
	p.context.removePlayer(p)
}

func (p *readerPlayer) IsPlaying() bool {
	p.m.Lock()
	defer p.m.Unlock()
	if err := p.ensurePlayer(); err != nil {
		p.context.setError(err)
		return false
	}

	return p.player.IsPlaying()
}

func (p *readerPlayer) Volume() float64 {
	p.m.Lock()
	defer p.m.Unlock()
	if err := p.ensurePlayer(); err != nil {
		p.context.setError(err)
		return 0
	}

	return p.player.Volume()
}

func (p *readerPlayer) SetVolume(volume float64) {
	p.m.Lock()
	defer p.m.Unlock()
	if err := p.ensurePlayer(); err != nil {
		p.context.setError(err)
		return
	}

	p.player.SetVolume(volume)
}

func (p *readerPlayer) Close() error {
	p.m.Lock()
	defer p.m.Unlock()
	runtime.SetFinalizer(p, nil)

	if p.player != nil {
		defer func() {
			p.player = nil
		}()
		p.player.Pause()
		return p.player.Close()
	}
	return nil
}

func (p *readerPlayer) Current() time.Duration {
	p.m.Lock()
	defer p.m.Unlock()
	if err := p.ensurePlayer(); err != nil {
		p.context.setError(err)
		return 0
	}

	sample := (p.stream.Current() - p.player.UnplayedBufferSize()) / bytesPerSample
	return time.Duration(sample) * time.Second / time.Duration(p.factory.sampleRate)
}

func (p *readerPlayer) Rewind() error {
	return p.Seek(0)
}

func (p *readerPlayer) Seek(offset time.Duration) error {
	p.m.Lock()
	defer p.m.Unlock()

	if p.player != nil {
		if p.player.IsPlaying() {
			defer func() {
				p.player.Play()
			}()
		}
		p.player.Reset()
	}
	return p.stream.Seek(offset)
}

func (p *readerPlayer) source() io.Reader {
	return p.stream.r
}

type timeStream struct {
	r          io.Reader
	sampleRate int
	pos        int64
	buf        []byte
	unread     int
}

func newTimeStream(r io.Reader, sampleRate int) (*timeStream, error) {
	s := &timeStream{
		r:          r,
		sampleRate: sampleRate,
	}
	if seeker, ok := s.r.(io.Seeker); ok {
		// Get the current position of the source.
		pos, err := seeker.Seek(0, io.SeekCurrent)
		if err != nil {
			return nil, err
		}
		s.pos = pos
	}
	return s, nil
}

func (s *timeStream) Unread(n int) {
	if s.unread+n > len(s.buf) {
		panic(fmt.Sprintf("audio: too much unreading: %d, the buffer size: %d, unreading position: %d", n, len(s.buf), s.unread))
	}
	s.unread += n
	s.pos -= int64(n)
}

func (s *timeStream) Read(buf []byte) (int, error) {
	if s.unread > 0 {
		n := copy(buf, s.buf[len(s.buf)-s.unread:])
		s.unread -= n
		s.pos += int64(n)
		return n, nil
	}

	n, err := s.r.Read(buf)
	s.pos += int64(n)
	s.buf = append(s.buf, buf[:n]...)
	if m := maxBufferSize(s.sampleRate); len(s.buf) > m {
		s.buf = s.buf[len(s.buf)-m:]
	}
	return n, err
}

func (s *timeStream) Seek(offset time.Duration) error {
	o := int64(offset) * bytesPerSample * int64(s.sampleRate) / int64(time.Second)

	// Align the byte position with the samples.
	o -= o % bytesPerSample
	o += s.pos % bytesPerSample

	seeker, ok := s.r.(io.Seeker)
	if !ok {
		panic("audio: the source must be io.Seeker when seeking but not")
	}
	pos, err := seeker.Seek(o, io.SeekStart)
	if err != nil {
		return err
	}

	s.pos = pos
	s.buf = s.buf[:0]
	s.unread = 0
	return nil
}

func (s *timeStream) Current() int64 {
	return s.pos
}
<|endoftext|>


<commit_before>// Copyright 2016 Apcera Inc. All rights reserved.

package server

import (
	"fmt"
	"sync"
	"runtime"
	"testing"

	"github.com/nats-io/nuid"
)

func createClientStore() *clientStore {
	cs := &clientStore{
		clients: make(map[string]*client),
	}
	return cs
}

func createClientInfo() (string, string) {
	nuid := nuid.New()

	clientID := "me"
	hbInbox := nuid.Next()

	return clientID, hbInbox
}

func TestClientRegister(t *testing.T) {
	cs := createClientStore()

	clientID, hbInbox := createClientInfo()

	// Register a new one
	c, isNew := cs.Register(clientID, hbInbox)
	if c == nil || !isNew {
		t.Fatal("Expected client to be new")
	}
	// Verify it's in the list of clients
	func() {
		cs.RLock()
		defer cs.RUnlock()

		if cs.clients[clientID] == nil {
			t.Fatal("Expected client to be registered")
		}
	}()
	// Verify the created client
	func() {
		c.RLock()
		defer c.RUnlock()
		if c.clientID != clientID {
			t.Fatalf("Expected client id to be %v, got %v", clientID, c.clientID)
		}
		if c.hbInbox != hbInbox {
			t.Fatalf("Expected client hbInbox to be %v, got %v", hbInbox, c.hbInbox)
		}
		if c.hbt != nil {
			t.Fatal("Did not expect timer to be set")
		}
		if c.fhb != 0 {
			t.Fatalf("Expected fhb to be 0, got %v", c.fhb)
		}
		if len(c.subs) != 0 {
			t.Fatalf("Expected subs count to be 0, got %v", len(c.subs))
		}
	}()

	// Register with same info
	secondCli, isNew := cs.Register(clientID, hbInbox)
	if secondCli != c || isNew {
		t.Fatal("Expected to get the same client")
	}
}

func TestClientParallelRegister(t *testing.T) {
	cs := createClientStore()

	_, hbInbox := createClientInfo()

	var wg sync.WaitGroup
	wg.Add(2)

	totalClients := 100

	for i := 0; i < 2; i++ {
		go func() {
			for j := 0; j < totalClients; j++ {
				clientID := fmt.Sprintf("clientID-%v", j)
				cs.Register(clientID, hbInbox)
				runtime.Gosched()
			}
			wg.Done()
		}()
	}

	wg.Wait()

	// We should not get more than totalClients
	func() {
		cs.RLock()
		defer cs.RUnlock()

		if len(cs.clients) != totalClients {
			t.Fatalf("Expected %v clients, got %v", totalClients, len(cs.clients))
		}
	}()
}

func TestClientUnregister(t *testing.T) {
	cs := createClientStore()

	clientID, hbInbox := createClientInfo()

	// Unregister one that does not exist should not cause a crash
	cs.Unregister(clientID)

	// Now register a client
	cs.Register(clientID, hbInbox)

	// Verify it's in the list of clients
	func() {
		cs.RLock()
		defer cs.RUnlock()

		if cs.clients[clientID] == nil {
			t.Fatal("Expected client to be registered")
		}
	}()

	// Unregister now
	cs.Unregister(clientID)

	// Verify it's gone.
	func() {
		cs.RLock()
		defer cs.RUnlock()

		if cs.clients[clientID] != nil {
			t.Fatal("Expected client to be unregistered")
		}
	}()
}

func TestClientLookup(t *testing.T) {
	cs := createClientStore()

	clientID, hbInbox := createClientInfo()

	// Looks-up one that does not exist
	if c := cs.Lookup("not-registered"); c != nil {
		t.Fatalf("Got unexpected client: %v", c)
	}

	// Registers one
	cs.Register(clientID, hbInbox)

	// Lookup again
	if c := cs.Lookup(clientID); c == nil {
		t.Fatal("Should have looked-up the client")
	}

	// Unregister
	cs.Unregister(clientID)

	// Lookup, should not be there
	if c := cs.Lookup(clientID); c != nil {
		t.Fatalf("Got unexpected client: %v", c)
	}
}

func TestClientGetClients(t *testing.T) {
	cs := createClientStore()

	clients := cs.GetClients()
	if len(clients) != 0 {
		t.Fatalf("Expected no client, got %v", len(clients))
	}

	nuid := nuid.New()

	clientID := "me"
	hbInbox := nuid.Next()

	cs.Register(clientID, hbInbox)

	clientID = "me2"
	hbInbox = nuid.Next()

	cs.Register(clientID, hbInbox)

	clients = cs.GetClients()
	if clients == nil || len(clients) != 2 {
		t.Fatalf("Expected to get 2 clients, got %v", len(clients))
	}

	for _, c := range clients {
		func() {
			c.RLock()
			defer c.RUnlock()

			if c.clientID != "me" && c.clientID != "me2" {
				t.Fatalf("Unexpected client ID: %v", c.clientID)
			}
		}()
	}
}

func TestClientAddSub(t *testing.T) {
	cs := createClientStore()

	clientID, hbInbox := createClientInfo()

	sub := &subState{}

	// Try to add a sub with client ID not registered
	if c := cs.AddSub(clientID, sub); c != nil {
		t.Fatalf("Expected AddSub to return nil, got %v", c)
	}

	// Now register the client
	cs.Register(clientID, hbInbox)

	// Now this should work
	c := cs.AddSub(clientID, sub)
	if c == nil {
		t.Fatal("Expected AddSub to return c")
	}

	// Check the sub is properly added to the client's subs list.
	func() {
		c.RLock()
		defer c.RUnlock()

		if len(c.subs) != 1 {
			t.Fatalf("Expected to have 1 sub, got %v", len(c.subs))
		}
		if c.subs[0] != sub {
			t.Fatalf("Got unexpected sub: %v", c.subs[0])
		}
	}()

	// Unregister
	cs.Unregister(clientID)

	// Again, this should fail since the clientID is not registered
	// anymore.
	if c := cs.AddSub(clientID, sub); c != nil {
		t.Fatalf("Expected AddSub to return nil, got %v", c)
	}
}

func TestClientRemoveSub(t *testing.T) {
	cs := createClientStore()

	clientID, hbInbox := createClientInfo()

	sub := &subState{}

	// Try to remove a sub with client ID not registered
	if c := cs.RemoveSub(clientID, sub); c != nil {
		t.Fatalf("Expected RemoveSub to return nil, got %v", c)
	}

	// Now register the client
	cs.Register(clientID, hbInbox)

	// Add a subscription
	c := cs.AddSub(clientID, sub)
	if c == nil {
		t.Fatal("Expected AddSub to return c")
	}

	// And remove it..
	if c := cs.RemoveSub(clientID, sub); c == nil {
		t.Fatal("Expected RemoveSub to return c")
	}

	// Unregister
	cs.Unregister(clientID)

	// Again, this should fail since the clientID is not registered
	// anymore.
	if c := cs.RemoveSub(clientID, sub); c != nil {
		t.Fatalf("Expected Remove to return nil, got %v", c)
	}
}

func TestClientGetSubs(t *testing.T) {
	cs := createClientStore()

	clientID, hbInbox := createClientInfo()

	if subs := cs.GetSubs(clientID); len(subs) != 0 {
		t.Fatalf("Expected 0 subs, got: %v", len(subs))
	}

	// Now register the client
	cs.Register(clientID, hbInbox)

	// Add a subscription
	c := cs.AddSub(clientID, &subState{subject: "foo"})
	if c == nil {
		t.Fatal("Expected AddSub to return c")
	}

	// or two
	// Add a subscription
	if c := cs.AddSub(clientID, &subState{subject: "bar"}); c == nil {
		t.Fatal("Expected AddSub to return c")
	}

	subs := cs.GetSubs(clientID)
	if len(subs) != 2 {
		t.Fatalf("Expected 2 subs, got: %v", len(subs))
	}

	for _, s := range subs {
		if s.subject != "foo" && s.subject != "bar" {
			t.Fatalf("Unexpected subject: %v", s.subject)
		}
	}
}
<commit_msg>Updated test based on comments<commit_after>// Copyright 2016 Apcera Inc. All rights reserved.

package server

import (
	"fmt"
	"runtime"
	"sync"
	"testing"

	"github.com/nats-io/nuid"
)

func createClientStore() *clientStore {
	cs := &clientStore{
		clients: make(map[string]*client),
	}
	return cs
}

func createClientInfo() (string, string) {
	nuid := nuid.New()

	clientID := "me"
	hbInbox := nuid.Next()

	return clientID, hbInbox
}

func TestClientRegister(t *testing.T) {
	cs := createClientStore()

	clientID, hbInbox := createClientInfo()

	// Register a new one
	c, isNew := cs.Register(clientID, hbInbox)
	if c == nil || !isNew {
		t.Fatal("Expected client to be new")
	}
	// Verify it's in the list of clients
	func() {
		cs.RLock()
		defer cs.RUnlock()

		if cs.clients[clientID] == nil {
			t.Fatal("Expected client to be registered")
		}
	}()
	// Verify the created client
	func() {
		c.RLock()
		defer c.RUnlock()
		if c.clientID != clientID {
			t.Fatalf("Expected client id to be %v, got %v", clientID, c.clientID)
		}
		if c.hbInbox != hbInbox {
			t.Fatalf("Expected client hbInbox to be %v, got %v", hbInbox, c.hbInbox)
		}
		if c.hbt != nil {
			t.Fatal("Did not expect timer to be set")
		}
		if c.fhb != 0 {
			t.Fatalf("Expected fhb to be 0, got %v", c.fhb)
		}
		if len(c.subs) != 0 {
			t.Fatalf("Expected subs count to be 0, got %v", len(c.subs))
		}
	}()

	// Register with same info
	secondCli, isNew := cs.Register(clientID, hbInbox)
	if secondCli != c || isNew {
		t.Fatal("Expected to get the same client")
	}
}

func TestClientParallelRegister(t *testing.T) {
	cs := createClientStore()

	_, hbInbox := createClientInfo()

	var wg sync.WaitGroup
	wg.Add(2)

	totalClients := 100
	errors := make(chan error, 2)

	for i := 0; i < 2; i++ {
		go func() {
			defer wg.Done()

			for j := 0; j < totalClients; j++ {
				clientID := fmt.Sprintf("clientID-%v", j)
				c, isNew := cs.Register(clientID, hbInbox)
				if c == nil {
					errors <- fmt.Errorf("client should not be nil")
					return
				}
				if !isNew && cs.Lookup(clientID) == nil {
					errors <- fmt.Errorf("Register returned isNew false, but clientID %v can't be found", clientID)
				}
				runtime.Gosched()
			}

		}()
	}

	wg.Wait()

	// Fail with the first error found.
	select {
	case e := <-errors:
		t.Fatalf("%v", e)
	default:
	}

	// We should not get more than totalClients
	func() {
		cs.RLock()
		defer cs.RUnlock()

		if len(cs.clients) != totalClients {
			t.Fatalf("Expected %v clients, got %v", totalClients, len(cs.clients))
		}
	}()
}

func TestClientUnregister(t *testing.T) {
	cs := createClientStore()

	clientID, hbInbox := createClientInfo()

	// Unregister one that does not exist should not cause a crash
	cs.Unregister(clientID)

	// Now register a client
	cs.Register(clientID, hbInbox)

	// Verify it's in the list of clients
	func() {
		cs.RLock()
		defer cs.RUnlock()

		if cs.clients[clientID] == nil {
			t.Fatal("Expected client to be registered")
		}
	}()

	// Unregister now
	cs.Unregister(clientID)

	// Verify it's gone.
	func() {
		cs.RLock()
		defer cs.RUnlock()

		if cs.clients[clientID] != nil {
			t.Fatal("Expected client to be unregistered")
		}
	}()
}

func TestClientLookup(t *testing.T) {
	cs := createClientStore()

	clientID, hbInbox := createClientInfo()

	// Looks-up one that does not exist
	if c := cs.Lookup("not-registered"); c != nil {
		t.Fatalf("Got unexpected client: %v", c)
	}

	// Registers one
	cs.Register(clientID, hbInbox)

	// Lookup again
	if c := cs.Lookup(clientID); c == nil {
		t.Fatal("Should have looked-up the client")
	}

	// Unregister
	cs.Unregister(clientID)

	// Lookup, should not be there
	if c := cs.Lookup(clientID); c != nil {
		t.Fatalf("Got unexpected client: %v", c)
	}
}

func TestClientGetClients(t *testing.T) {
	cs := createClientStore()

	clients := cs.GetClients()
	if len(clients) != 0 {
		t.Fatalf("Expected no client, got %v", len(clients))
	}

	nuid := nuid.New()

	clientID := "me"
	hbInbox := nuid.Next()

	cs.Register(clientID, hbInbox)

	clientID = "me2"
	hbInbox = nuid.Next()

	cs.Register(clientID, hbInbox)

	clients = cs.GetClients()
	if clients == nil || len(clients) != 2 {
		t.Fatalf("Expected to get 2 clients, got %v", len(clients))
	}

	for _, c := range clients {
		func() {
			c.RLock()
			defer c.RUnlock()

			if c.clientID != "me" && c.clientID != "me2" {
				t.Fatalf("Unexpected client ID: %v", c.clientID)
			}
		}()
	}
}

func TestClientAddSub(t *testing.T) {
	cs := createClientStore()

	clientID, hbInbox := createClientInfo()

	sub := &subState{}

	// Try to add a sub with client ID not registered
	if c := cs.AddSub(clientID, sub); c != nil {
		t.Fatalf("Expected AddSub to return nil, got %v", c)
	}

	// Now register the client
	cs.Register(clientID, hbInbox)

	// Now this should work
	c := cs.AddSub(clientID, sub)
	if c == nil {
		t.Fatal("Expected AddSub to return c")
	}

	// Check the sub is properly added to the client's subs list.
	func() {
		c.RLock()
		defer c.RUnlock()

		if len(c.subs) != 1 {
			t.Fatalf("Expected to have 1 sub, got %v", len(c.subs))
		}
		if c.subs[0] != sub {
			t.Fatalf("Got unexpected sub: %v", c.subs[0])
		}
	}()

	// Unregister
	cs.Unregister(clientID)

	// Again, this should fail since the clientID is not registered
	// anymore.
	if c := cs.AddSub(clientID, sub); c != nil {
		t.Fatalf("Expected AddSub to return nil, got %v", c)
	}
}

func TestClientRemoveSub(t *testing.T) {
	cs := createClientStore()

	clientID, hbInbox := createClientInfo()

	sub := &subState{}

	// Try to remove a sub with client ID not registered
	if c := cs.RemoveSub(clientID, sub); c != nil {
		t.Fatalf("Expected RemoveSub to return nil, got %v", c)
	}

	// Now register the client
	cs.Register(clientID, hbInbox)

	// Add a subscription
	c := cs.AddSub(clientID, sub)
	if c == nil {
		t.Fatal("Expected AddSub to return c")
	}

	// And remove it..
	if c := cs.RemoveSub(clientID, sub); c == nil {
		t.Fatal("Expected RemoveSub to return c")
	}

	// Unregister
	cs.Unregister(clientID)

	// Again, this should fail since the clientID is not registered
	// anymore.
	if c := cs.RemoveSub(clientID, sub); c != nil {
		t.Fatalf("Expected Remove to return nil, got %v", c)
	}
}

func TestClientGetSubs(t *testing.T) {
	cs := createClientStore()

	clientID, hbInbox := createClientInfo()

	if subs := cs.GetSubs(clientID); len(subs) != 0 {
		t.Fatalf("Expected 0 subs, got: %v", len(subs))
	}

	// Now register the client
	cs.Register(clientID, hbInbox)

	// Add a subscription
	c := cs.AddSub(clientID, &subState{subject: "foo"})
	if c == nil {
		t.Fatal("Expected AddSub to return c")
	}

	// or two
	// Add a subscription
	if c := cs.AddSub(clientID, &subState{subject: "bar"}); c == nil {
		t.Fatal("Expected AddSub to return c")
	}

	subs := cs.GetSubs(clientID)
	if len(subs) != 2 {
		t.Fatalf("Expected 2 subs, got: %v", len(subs))
	}

	for _, s := range subs {
		if s.subject != "foo" && s.subject != "bar" {
			t.Fatalf("Unexpected subject: %v", s.subject)
		}
	}
}
<|endoftext|>


<commit_before>package heroku

import (
	"fmt"
	"net/http"
	"time"

	"github.com/remind101/empire"
	streamhttp "github.com/remind101/empire/pkg/stream/http"
	"golang.org/x/net/context"
)

type PostLogs struct {
	*empire.Empire
}

type PostLogsForm struct {
	Duration int64
}

func (h *PostLogs) ServeHTTPContext(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
	a, err := findApp(ctx, h)
	if err != nil {
		return fmt.Errorf("error finding app: %v", err)
	}

	var form PostLogsForm
	if err := Decode(r, &form); err != nil {
		if err.Error() != "EOF" {
			return fmt.Errorf("error decoding request: %v", err)
		}
	}

	rw := streamhttp.StreamingResponseWriter(w)

	// Prevent the ELB idle connection timeout to close the connection.
	defer close(streamhttp.Heartbeat(rw, 10*time.Second))

	err = h.StreamLogs(a, rw, time.Duration(form.Duration))
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>Don't annotate the error since we handle the gorm exception upstream<commit_after>package heroku

import (
	"fmt"
	"net/http"
	"time"

	"github.com/remind101/empire"
	streamhttp "github.com/remind101/empire/pkg/stream/http"
	"golang.org/x/net/context"
)

type PostLogs struct {
	*empire.Empire
}

type PostLogsForm struct {
	Duration int64
}

func (h *PostLogs) ServeHTTPContext(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
	a, err := findApp(ctx, h)
	if err != nil {
		return err
	}

	var form PostLogsForm
	if err := Decode(r, &form); err != nil {
		if err.Error() != "EOF" {
			return fmt.Errorf("error decoding request: %v", err)
		}
	}

	rw := streamhttp.StreamingResponseWriter(w)

	// Prevent the ELB idle connection timeout to close the connection.
	defer close(streamhttp.Heartbeat(rw, 10*time.Second))

	err = h.StreamLogs(a, rw, time.Duration(form.Duration))
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package yaml

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"strings"

	"github.com/qor/i18n"
	"gopkg.in/yaml.v2"
)

// New new YAML backend for I18n
func New(paths ...string) i18n.Backend {
	backend := &Backend{}

	var files []string
	for _, p := range paths {
		if file, err := os.Open(p); err == nil {
			defer file.Close()
			if fileInfo, err := file.Stat(); err == nil {
				if fileInfo.IsDir() {
					yamlFiles, _ := filepath.Glob(path.Join(p, "*.yaml"))
					files = append(files, yamlFiles...)

					ymlFiles, _ := filepath.Glob(path.Join(p, "*.yml"))
					files = append(files, ymlFiles...)
				} else if fileInfo.Mode().IsRegular() {
					files = append(files, p)
				}
			}
		}
	}
	for _, file := range files {
		if content, err := ioutil.ReadFile(file); err == nil {
			backend.contents = append(backend.contents, content)
		}
	}
	return backend
}

// NewWithWalk has the same functionality as New but uses filepath.Walk to find all the translation files recursively.
func NewWithWalk(paths ...string) i18n.Backend {
	backend := &Backend{}

	var files []string
	for _, p := range paths {
		filepath.Walk(p, func(path string, fileInfo os.FileInfo, err error) error {
			if isYamlFile(fileInfo) {
				files = append(files, path)
			}
			return nil
		})
	}
	for _, file := range files {
		if content, err := ioutil.ReadFile(file); err == nil {
			backend.contents = append(backend.contents, content)
		}
	}

	return backend
}

func isYamlFile(fileInfo os.FileInfo) bool {
	if fileInfo == nil {
		return false
	}
	return fileInfo.Mode().IsRegular() && (strings.HasSuffix(fileInfo.Name(), ".yml") || strings.HasSuffix(fileInfo.Name(), ".yaml"))
}

func walkFilesystem(fs http.FileSystem, entry http.File, prefix string) [][]byte {
	var (
		contents [][]byte
		err      error
		isRoot   bool
	)
	if entry == nil {
		if entry, err = fs.Open("/"); err != nil {
			return nil
		}
		isRoot = true
		defer entry.Close()
	}
	fileInfo, err := entry.Stat()
	if err != nil {
		return nil
	}
	if !isRoot {
		prefix = prefix + fileInfo.Name() + "/"
	}
	if fileInfo.IsDir() {
		if entries, err := entry.Readdir(-1); err == nil {
			for _, e := range entries {
				if file, err := fs.Open(prefix + e.Name()); err == nil {
					defer file.Close()
					contents = append(contents, walkFilesystem(fs, file, prefix)...)
				}
			}
		}
	} else if isYamlFile(fileInfo) {
		if content, err := ioutil.ReadAll(entry); err == nil {
			contents = append(contents, content)
		}
	}
	return contents
}

// NewWithFilesystem initializes a backend that reads translation files from an http.FileSystem.
func NewWithFilesystem(fss ...http.FileSystem) i18n.Backend {
	backend := &Backend{}

	for _, fs := range fss {
		backend.contents = append(backend.contents, walkFilesystem(fs, nil, "/")...)
	}
	return backend
}

// Backend YAML backend
type Backend struct {
	contents [][]byte
}

func loadTranslationsFromYaml(locale string, value interface{}, scopes []string) (translations []*i18n.Translation) {
	switch v := value.(type) {
	case yaml.MapSlice:
		for _, s := range v {
			results := loadTranslationsFromYaml(locale, s.Value, append(scopes, fmt.Sprint(s.Key)))
			translations = append(translations, results...)
		}
	default:
		var translation = &i18n.Translation{
			Locale: locale,
			Key:    strings.Join(scopes, "."),
			Value:  fmt.Sprint(v),
		}
		translations = append(translations, translation)
	}
	return
}

// LoadYAMLContent load YAML content
func (backend *Backend) LoadYAMLContent(content []byte) (translations []*i18n.Translation, err error) {
	var slice yaml.MapSlice

	if err = yaml.Unmarshal(content, &slice); err == nil {
		for _, item := range slice {
			translations = append(translations, loadTranslationsFromYaml(item.Key.(string) /* locale */, item.Value, []string{})...)
		}
	}

	return translations, err
}

// LoadTranslations load translations from YAML backend
func (backend *Backend) LoadTranslations() (translations []*i18n.Translation) {
	for _, content := range backend.contents {
		if results, err := backend.LoadYAMLContent(content); err == nil {
			translations = append(translations, results...)
		} else {
			panic(err)
		}
	}
	return translations
}

// SaveTranslation save translation into YAML backend, not implemented
func (backend *Backend) SaveTranslation(t *i18n.Translation) error {
	return errors.New("not implemented")
}

// DeleteTranslation delete translation into YAML backend, not implemented
func (backend *Backend) DeleteTranslation(t *i18n.Translation) error {
	return errors.New("not implemented")
}
<commit_msg>YAML return itself other than Backend interface<commit_after>package yaml

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"strings"

	"github.com/qor/i18n"
	"gopkg.in/yaml.v2"
)

var _ i18n.Backend = &Backend{}

// New new YAML backend for I18n
func New(paths ...string) *Backend {
	backend := &Backend{}

	var files []string
	for _, p := range paths {
		if file, err := os.Open(p); err == nil {
			defer file.Close()
			if fileInfo, err := file.Stat(); err == nil {
				if fileInfo.IsDir() {
					yamlFiles, _ := filepath.Glob(path.Join(p, "*.yaml"))
					files = append(files, yamlFiles...)

					ymlFiles, _ := filepath.Glob(path.Join(p, "*.yml"))
					files = append(files, ymlFiles...)
				} else if fileInfo.Mode().IsRegular() {
					files = append(files, p)
				}
			}
		}
	}
	for _, file := range files {
		if content, err := ioutil.ReadFile(file); err == nil {
			backend.contents = append(backend.contents, content)
		}
	}
	return backend
}

// NewWithWalk has the same functionality as New but uses filepath.Walk to find all the translation files recursively.
func NewWithWalk(paths ...string) i18n.Backend {
	backend := &Backend{}

	var files []string
	for _, p := range paths {
		filepath.Walk(p, func(path string, fileInfo os.FileInfo, err error) error {
			if isYamlFile(fileInfo) {
				files = append(files, path)
			}
			return nil
		})
	}
	for _, file := range files {
		if content, err := ioutil.ReadFile(file); err == nil {
			backend.contents = append(backend.contents, content)
		}
	}

	return backend
}

func isYamlFile(fileInfo os.FileInfo) bool {
	if fileInfo == nil {
		return false
	}
	return fileInfo.Mode().IsRegular() && (strings.HasSuffix(fileInfo.Name(), ".yml") || strings.HasSuffix(fileInfo.Name(), ".yaml"))
}

func walkFilesystem(fs http.FileSystem, entry http.File, prefix string) [][]byte {
	var (
		contents [][]byte
		err      error
		isRoot   bool
	)
	if entry == nil {
		if entry, err = fs.Open("/"); err != nil {
			return nil
		}
		isRoot = true
		defer entry.Close()
	}
	fileInfo, err := entry.Stat()
	if err != nil {
		return nil
	}
	if !isRoot {
		prefix = prefix + fileInfo.Name() + "/"
	}
	if fileInfo.IsDir() {
		if entries, err := entry.Readdir(-1); err == nil {
			for _, e := range entries {
				if file, err := fs.Open(prefix + e.Name()); err == nil {
					defer file.Close()
					contents = append(contents, walkFilesystem(fs, file, prefix)...)
				}
			}
		}
	} else if isYamlFile(fileInfo) {
		if content, err := ioutil.ReadAll(entry); err == nil {
			contents = append(contents, content)
		}
	}
	return contents
}

// NewWithFilesystem initializes a backend that reads translation files from an http.FileSystem.
func NewWithFilesystem(fss ...http.FileSystem) i18n.Backend {
	backend := &Backend{}

	for _, fs := range fss {
		backend.contents = append(backend.contents, walkFilesystem(fs, nil, "/")...)
	}
	return backend
}

// Backend YAML backend
type Backend struct {
	contents [][]byte
}

func loadTranslationsFromYaml(locale string, value interface{}, scopes []string) (translations []*i18n.Translation) {
	switch v := value.(type) {
	case yaml.MapSlice:
		for _, s := range v {
			results := loadTranslationsFromYaml(locale, s.Value, append(scopes, fmt.Sprint(s.Key)))
			translations = append(translations, results...)
		}
	default:
		var translation = &i18n.Translation{
			Locale: locale,
			Key:    strings.Join(scopes, "."),
			Value:  fmt.Sprint(v),
		}
		translations = append(translations, translation)
	}
	return
}

// LoadYAMLContent load YAML content
func (backend *Backend) LoadYAMLContent(content []byte) (translations []*i18n.Translation, err error) {
	var slice yaml.MapSlice

	if err = yaml.Unmarshal(content, &slice); err == nil {
		for _, item := range slice {
			translations = append(translations, loadTranslationsFromYaml(item.Key.(string) /* locale */, item.Value, []string{})...)
		}
	}

	return translations, err
}

// LoadTranslations load translations from YAML backend
func (backend *Backend) LoadTranslations() (translations []*i18n.Translation) {
	for _, content := range backend.contents {
		if results, err := backend.LoadYAMLContent(content); err == nil {
			translations = append(translations, results...)
		} else {
			panic(err)
		}
	}
	return translations
}

// SaveTranslation save translation into YAML backend, not implemented
func (backend *Backend) SaveTranslation(t *i18n.Translation) error {
	return errors.New("not implemented")
}

// DeleteTranslation delete translation into YAML backend, not implemented
func (backend *Backend) DeleteTranslation(t *i18n.Translation) error {
	return errors.New("not implemented")
}
<|endoftext|>