<commit_before>package flv

import (
    "fmt"
    "bytes"
)

type AVCProfile byte
const (
    AVC_PROFILE_BASELINE   AVCProfile = 66
    AVC_PROFILE_MAIN       AVCProfile = 77
    AVC_PROFILE_EXTENDED   AVCProfile = 88
    AVC_PROFILE_HIGH       AVCProfile = 100
    AVC_PROFILE_HIGH10     AVCProfile = 110
    AVC_PROFILE_HIGH422    AVCProfile = 122
    AVC_PROFILE_HIGH444    AVCProfile = 244
    AVC_PROFILE_CAVLC444   AVCProfile = 44
)

var (
    avcProfileStrings = map[AVCProfile]string {
        AVC_PROFILE_BASELINE:   "Baseline",
        AVC_PROFILE_MAIN:       "Main",
        AVC_PROFILE_EXTENDED:   "Extended",
        AVC_PROFILE_HIGH:       "High",
        AVC_PROFILE_HIGH10:     "High 10",
        AVC_PROFILE_HIGH422:    "High 4:2:2",
        AVC_PROFILE_HIGH444:    "High 4:4:4",
        AVC_PROFILE_CAVLC444:   "CAVLC 4:4:4",
    }
)

func (p AVCProfile) String() string {
    return avcProfileStrings[p]
}

type AVCConfRecord struct {
    ConfigurationVersion byte
    AVCProfileIndication AVCProfile
    ProfileCompatibility byte
    AVCLevelIndication  byte
    RawSPSData [][]byte
    RawPPSData [][]byte
}

func (r *AVCConfRecord) String() string {
    return fmt.Sprintf("AVCConfigurationRecord(ver. %d, profile: %s, level: %d, %d SPS, %d PPS)",
        r.ConfigurationVersion, r.AVCProfileIndication,
        r.AVCLevelIndication,
        len(r.RawSPSData), len(r.RawPPSData))
}

func ParseAVCConfRecord(data []byte) (rec *AVCConfRecord, err error) {
    r := NewBitReader(data)

    defer func () {
        if rec := recover(); rec != nil {
            err = rec.(error)
        }
    }()

    configurationVersion := r.U8()
    AVCProfileIndication := r.U8()
    profile_compatibility := r.U8()
    AVCLevelIndication := r.U8()

    if r.U(6) != 077 {
        panic("wrong reserved 1")
    }

    r.U(2) /* lengthSizeMinusOne */
    if r.U(3) != 07 {
        panic("wrong reserved 2")
    }

    numOfSPS := r.U(5)
    spss := make([][]byte, numOfSPS)
    for i := uint32(0); i < numOfSPS; i++ {
        spsLen := r.U(16)
        spss[i] = make([]byte, spsLen)
        r.Read(spss[i])
    }

    numOfPPS := r.U(8)
    ppss := make([][]byte, numOfPPS)
    for i := uint32(0); i < numOfPPS; i++ {
        ppsLen := r.U(16)
        ppss[i] = make([]byte, ppsLen)
        r.Read(ppss[i])
    }

    rec = &AVCConfRecord{
            ConfigurationVersion: configurationVersion,
            AVCProfileIndication: AVCProfile(AVCProfileIndication),
            ProfileCompatibility: profile_compatibility,
            AVCLevelIndication: AVCLevelIndication,
            RawSPSData: spss,
            RawPPSData: ppss,
        }
    return
}

type SPS struct {
    Profile_idc AVCProfile
    Constraint_set byte
    Level_idc byte
    SPS_id uint32

    pic_width_in_mbs uint32
    pic_height_in_map_units uint32
    frame_mbs_only_flag uint32
    crops FrameCropOffsets
}

type FrameCropOffsets struct {
    left uint32
    right uint32
    top uint32
    bottom uint32
}

func (sps *SPS) Width() uint32 {
    w := sps.pic_width_in_mbs*16 - sps.crops.left*2 - sps.crops.right*2
    return w
}

func (sps *SPS) Height() uint32 {
    c := uint32(2) - sps.frame_mbs_only_flag
    h := sps.pic_height_in_map_units*16 - sps.crops.top*2 - sps.crops.bottom*2
    return c*h
}

func (sps *SPS) String() string {
    return fmt.Sprintf("seq_parameter_set(profile: %d, level: %d, id: %d)", sps.Profile_idc, sps.Level_idc, sps.SPS_id)
}

func ParseSPS(rawSPSNALU []byte) (ret *SPS, err error) {
    r := NewBitReader(rawSPSNALU)

    defer func () {
        if rec := recover(); rec != nil {
            err = rec.(error)
        }
    }()

    r.U(1) /* forbidden_zero_bit */
    r.U(2) /* nal_ref_idc */

    nal_unit_type := r.U(5)
    if nal_unit_type != 7 {
        err = fmt.Errorf("Not SPS NALU, nal_unit_type = %d", nal_unit_type)
        return
    }

    profile_idc := r.U8()

    constraint_set_flags := byte(0)
    for i := uint(6); i > 0; i-- {
        f := byte(r.U(1))
        constraint_set_flags |= f << (i - 1)
    }

    r.U(2) /* reserved_zero_2bits */
    level_idc := r.U8()

    seq_parameter_set_id := r.Ue()

    extended_profiles := []byte{100, 110, 122, 244, 44, 83, 86, 118, 128}
    if bytes.IndexByte(extended_profiles, profile_idc) != -1 {

        chroma_format_idc := r.Ue()
        if chroma_format_idc == 3 {
            r.U(1) // separate_colour_plane_flag
        }
        r.Ue() // bit_depth_luma_minus8
        r.Ue() // bit_depth_chroma_minus8
        r.U(1) // qpprime_y_zero_transform_bypass_flag
        seq_scaling_matrix_present_flag := r.U(1)
        if seq_scaling_matrix_present_flag != 0 {
            c := 12
            if chroma_format_idc != 3 {
                c = 8
            }

            for i := 0; i < c; i++ {
                seq_scaling_list_present_flag := r.U(1)
                if seq_scaling_list_present_flag != 0 {
                    if i < 6 {
                        scaling_list(r, 16)
                    } else {
                        scaling_list(r, 64)
                    }
                }
            }
        }
    }

    r.Ue() /* log2_max_frame_num_minus4 */
    pic_order_cnt_type := r.Ue()
    if pic_order_cnt_type == 0 {
        r.Ue() /* log2_max_pic_order_cnt_lsb_minus4 */
    } else if pic_order_cnt_type == 1 {
        r.U(1) /* delta_pic_order_always_zero_flag */
        r.Se() /* offset_for_non_ref_pic */
        r.Se() /* offset_for_top_to_bottom_field */
        num_ref_frames_in_pic_order_cnt_cycle := r.Ue()
        for i := uint32(0); i <num_ref_frames_in_pic_order_cnt_cycle; i++ {
            r.Se() /* offset_for_ref_frame[ i ] */
        }
    }

    r.Ue() /* max_num_ref_frames */
    r.U(1) /* gaps_in_frame_num_value_allowed_flag */
    pic_width_in_mbs_minus1 := r.Ue()
    pic_height_in_map_units_minus1 := r.Ue()

    frame_mbs_only_flag := r.U(1)
    if frame_mbs_only_flag == 0 {
        r.U(1) /* mb_adaptive_frame_field_flag */
    }

    r.U(1) /* direct_8x8_inference_flag */

    crops := FrameCropOffsets{}

    frame_cropping_flag := r.U(1)
    if frame_cropping_flag != 0 {
        frame_crop_left_offset := r.Ue()
        frame_crop_right_offset := r.Ue()
        frame_crop_top_offset := r.Ue()
        frame_crop_bottom_offset := r.Ue()

        crops = FrameCropOffsets{
                    left: frame_crop_left_offset,
                    right: frame_crop_right_offset,
                    top: frame_crop_top_offset,
                    bottom: frame_crop_bottom_offset,
                }
    }

    ret = &SPS{
                Profile_idc: AVCProfile(profile_idc),
                Constraint_set: constraint_set_flags,
                Level_idc: level_idc,
                SPS_id: seq_parameter_set_id,

                pic_width_in_mbs: pic_width_in_mbs_minus1 + 1,
                pic_height_in_map_units : pic_height_in_map_units_minus1 + 1,
                frame_mbs_only_flag : frame_mbs_only_flag,
                crops: crops,
            }
    return
}

func scaling_list(r *BitReader, scalingListSize uint32) {
    lastScale := int32(8)
    nextScale := int32(8)

    for j := uint32(0); j < scalingListSize; j++ {
        if nextScale != 0 {
            delta_scale := r.Se()
            nextScale = (lastScale + delta_scale + 256) % 256
        }
        if nextScale != 0 {
            lastScale = nextScale
        }
    }
}<commit_msg>nobody follows standard =(<commit_after>package flv

import (
    "fmt"
    "bytes"
)

type AVCProfile byte
const (
    AVC_PROFILE_BASELINE   AVCProfile = 66
    AVC_PROFILE_MAIN       AVCProfile = 77
    AVC_PROFILE_EXTENDED   AVCProfile = 88
    AVC_PROFILE_HIGH       AVCProfile = 100
    AVC_PROFILE_HIGH10     AVCProfile = 110
    AVC_PROFILE_HIGH422    AVCProfile = 122
    AVC_PROFILE_HIGH444    AVCProfile = 244
    AVC_PROFILE_CAVLC444   AVCProfile = 44
)

var (
    avcProfileStrings = map[AVCProfile]string {
        AVC_PROFILE_BASELINE:   "Baseline",
        AVC_PROFILE_MAIN:       "Main",
        AVC_PROFILE_EXTENDED:   "Extended",
        AVC_PROFILE_HIGH:       "High",
        AVC_PROFILE_HIGH10:     "High 10",
        AVC_PROFILE_HIGH422:    "High 4:2:2",
        AVC_PROFILE_HIGH444:    "High 4:4:4",
        AVC_PROFILE_CAVLC444:   "CAVLC 4:4:4",
    }
)

func (p AVCProfile) String() string {
    return avcProfileStrings[p]
}

type AVCConfRecord struct {
    ConfigurationVersion byte
    AVCProfileIndication AVCProfile
    ProfileCompatibility byte
    AVCLevelIndication  byte
    RawSPSData [][]byte
    RawPPSData [][]byte
}

func (r *AVCConfRecord) String() string {
    return fmt.Sprintf("AVCConfigurationRecord(ver. %d, profile: %s, level: %d, %d SPS, %d PPS)",
        r.ConfigurationVersion, r.AVCProfileIndication,
        r.AVCLevelIndication,
        len(r.RawSPSData), len(r.RawPPSData))
}

func ParseAVCConfRecord(data []byte) (rec *AVCConfRecord, err error) {
    r := NewBitReader(data)

    defer func () {
        if rec := recover(); rec != nil {
            err = rec.(error)
        }
    }()

    configurationVersion := r.U8()
    AVCProfileIndication := r.U8()
    profile_compatibility := r.U8()
    AVCLevelIndication := r.U8()

    r.U(6)
    /* nobody follows standard in reserved
    if r.U(6) != 077 {
        panic("wrong reserved 1")
    } */

    r.U(2) /* lengthSizeMinusOne */
    r.U(3)
    /* same here
    if r.U(3) != 07 {
        panic("wrong reserved 2")
    } */

    numOfSPS := r.U(5)
    spss := make([][]byte, numOfSPS)
    for i := uint32(0); i < numOfSPS; i++ {
        spsLen := r.U(16)
        spss[i] = make([]byte, spsLen)
        r.Read(spss[i])
    }

    numOfPPS := r.U(8)
    ppss := make([][]byte, numOfPPS)
    for i := uint32(0); i < numOfPPS; i++ {
        ppsLen := r.U(16)
        ppss[i] = make([]byte, ppsLen)
        r.Read(ppss[i])
    }

    rec = &AVCConfRecord{
            ConfigurationVersion: configurationVersion,
            AVCProfileIndication: AVCProfile(AVCProfileIndication),
            ProfileCompatibility: profile_compatibility,
            AVCLevelIndication: AVCLevelIndication,
            RawSPSData: spss,
            RawPPSData: ppss,
        }
    return
}

type SPS struct {
    Profile_idc AVCProfile
    Constraint_set byte
    Level_idc byte
    SPS_id uint32

    pic_width_in_mbs uint32
    pic_height_in_map_units uint32
    frame_mbs_only_flag uint32
    crops FrameCropOffsets
}

type FrameCropOffsets struct {
    left uint32
    right uint32
    top uint32
    bottom uint32
}

func (sps *SPS) Width() uint32 {
    w := sps.pic_width_in_mbs*16 - sps.crops.left*2 - sps.crops.right*2
    return w
}

func (sps *SPS) Height() uint32 {
    c := uint32(2) - sps.frame_mbs_only_flag
    h := sps.pic_height_in_map_units*16 - sps.crops.top*2 - sps.crops.bottom*2
    return c*h
}

func (sps *SPS) String() string {
    return fmt.Sprintf("seq_parameter_set(profile: %d, level: %d, id: %d)", sps.Profile_idc, sps.Level_idc, sps.SPS_id)
}

func ParseSPS(rawSPSNALU []byte) (ret *SPS, err error) {
    r := NewBitReader(rawSPSNALU)

    defer func () {
        if rec := recover(); rec != nil {
            err = rec.(error)
        }
    }()

    r.U(1) /* forbidden_zero_bit */
    r.U(2) /* nal_ref_idc */

    nal_unit_type := r.U(5)
    if nal_unit_type != 7 {
        err = fmt.Errorf("Not SPS NALU, nal_unit_type = %d", nal_unit_type)
        return
    }

    profile_idc := r.U8()

    constraint_set_flags := byte(0)
    for i := uint(6); i > 0; i-- {
        f := byte(r.U(1))
        constraint_set_flags |= f << (i - 1)
    }

    r.U(2) /* reserved_zero_2bits */
    level_idc := r.U8()

    seq_parameter_set_id := r.Ue()

    extended_profiles := []byte{100, 110, 122, 244, 44, 83, 86, 118, 128}
    if bytes.IndexByte(extended_profiles, profile_idc) != -1 {

        chroma_format_idc := r.Ue()
        if chroma_format_idc == 3 {
            r.U(1) // separate_colour_plane_flag
        }
        r.Ue() // bit_depth_luma_minus8
        r.Ue() // bit_depth_chroma_minus8
        r.U(1) // qpprime_y_zero_transform_bypass_flag
        seq_scaling_matrix_present_flag := r.U(1)
        if seq_scaling_matrix_present_flag != 0 {
            c := 12
            if chroma_format_idc != 3 {
                c = 8
            }

            for i := 0; i < c; i++ {
                seq_scaling_list_present_flag := r.U(1)
                if seq_scaling_list_present_flag != 0 {
                    if i < 6 {
                        scaling_list(r, 16)
                    } else {
                        scaling_list(r, 64)
                    }
                }
            }
        }
    }

    r.Ue() /* log2_max_frame_num_minus4 */
    pic_order_cnt_type := r.Ue()
    if pic_order_cnt_type == 0 {
        r.Ue() /* log2_max_pic_order_cnt_lsb_minus4 */
    } else if pic_order_cnt_type == 1 {
        r.U(1) /* delta_pic_order_always_zero_flag */
        r.Se() /* offset_for_non_ref_pic */
        r.Se() /* offset_for_top_to_bottom_field */
        num_ref_frames_in_pic_order_cnt_cycle := r.Ue()
        for i := uint32(0); i <num_ref_frames_in_pic_order_cnt_cycle; i++ {
            r.Se() /* offset_for_ref_frame[ i ] */
        }
    }

    r.Ue() /* max_num_ref_frames */
    r.U(1) /* gaps_in_frame_num_value_allowed_flag */
    pic_width_in_mbs_minus1 := r.Ue()
    pic_height_in_map_units_minus1 := r.Ue()

    frame_mbs_only_flag := r.U(1)
    if frame_mbs_only_flag == 0 {
        r.U(1) /* mb_adaptive_frame_field_flag */
    }

    r.U(1) /* direct_8x8_inference_flag */

    crops := FrameCropOffsets{}

    frame_cropping_flag := r.U(1)
    if frame_cropping_flag != 0 {
        frame_crop_left_offset := r.Ue()
        frame_crop_right_offset := r.Ue()
        frame_crop_top_offset := r.Ue()
        frame_crop_bottom_offset := r.Ue()

        crops = FrameCropOffsets{
                    left: frame_crop_left_offset,
                    right: frame_crop_right_offset,
                    top: frame_crop_top_offset,
                    bottom: frame_crop_bottom_offset,
                }
    }

    ret = &SPS{
                Profile_idc: AVCProfile(profile_idc),
                Constraint_set: constraint_set_flags,
                Level_idc: level_idc,
                SPS_id: seq_parameter_set_id,

                pic_width_in_mbs: pic_width_in_mbs_minus1 + 1,
                pic_height_in_map_units : pic_height_in_map_units_minus1 + 1,
                frame_mbs_only_flag : frame_mbs_only_flag,
                crops: crops,
            }
    return
}

func scaling_list(r *BitReader, scalingListSize uint32) {
    lastScale := int32(8)
    nextScale := int32(8)

    for j := uint32(0); j < scalingListSize; j++ {
        if nextScale != 0 {
            delta_scale := r.Se()
            nextScale = (lastScale + delta_scale + 256) % 256
        }
        if nextScale != 0 {
            lastScale = nextScale
        }
    }
}
<|endoftext|>


<commit_before>package backend

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

	"code.google.com/p/go-uuid/uuid"

	"github.com/fsouza/go-dockerclient"
	"github.com/pkg/sftp"
	"github.com/travis-ci/worker/lib/context"
	"github.com/travis-ci/worker/lib/metrics"
	"golang.org/x/crypto/ssh"
	gocontext "golang.org/x/net/context"
)

type DockerProvider struct {
	client *docker.Client

	cpuSetsMutex sync.Mutex
	cpuSets      []bool
}

type DockerInstance struct {
	client    *docker.Client
	provider  *DockerProvider
	container *docker.Container
}

func NewDockerProvider(config map[string]string) (*DockerProvider, error) {
	endpoint, ok := config["endpoint"]
	if !ok {
		return nil, fmt.Errorf("expected config key endpoint")
	}

	client, err := docker.NewClient(endpoint)
	if err != nil {
		return nil, err
	}

	cpuSetSize := runtime.NumCPU()
	if cpuSetSize < 2 {
		cpuSetSize = 2
	}

	return &DockerProvider{
		client:  client,
		cpuSets: make([]bool, cpuSetSize),
	}, nil
}

func (p *DockerProvider) Start(ctx gocontext.Context, startAttributes StartAttributes) (Instance, error) {
	cpuSets, err := p.checkoutCPUSets()
	if err != nil && cpuSets != "" {
		return nil, err
	}

	imageID, err := p.imageForLanguage(startAttributes.Language)
	if err != nil {
		return nil, err
	}

	dockerConfig := &docker.Config{
		Cmd:      []string{"/sbin/init"},
		Image:    imageID,
		Memory:   1024 * 1024 * 1024 * 4,
		Hostname: fmt.Sprintf("testing-go-%s", uuid.NewUUID()),
	}
	if cpuSets != "" {
		dockerConfig.CPUSet = cpuSets
	}

	container, err := p.client.CreateContainer(docker.CreateContainerOptions{Config: dockerConfig})
	if err != nil {
		if container != nil {
			err := p.client.RemoveContainer(docker.RemoveContainerOptions{
				ID:            container.ID,
				RemoveVolumes: true,
				Force:         true,
			})
			if err != nil {
				context.LoggerFromContext(ctx).WithField("err", err).Error("couldn't remove container after create failure")
			}
		}

		return nil, err
	}

	startBooting := time.Now()

	err = p.client.StartContainer(container.ID, &docker.HostConfig{})
	if err != nil {
		return nil, err
	}

	containerReady := make(chan *docker.Container)
	errChan := make(chan error)
	go func(id string) {
		for true {
			container, err := p.client.InspectContainer(id)
			if err != nil {
				errChan <- err
				return
			}

			if container.State.Running {
				containerReady <- container
				return
			}
		}
	}(container.ID)

	select {
	case container := <-containerReady:
		metrics.TimeSince("worker.vm.provider.docker.boot", startBooting)
		return &DockerInstance{
			client:    p.client,
			provider:  p,
			container: container,
		}, nil
	case err := <-errChan:
		return nil, err
	case <-ctx.Done():
		if ctx.Err() == gocontext.DeadlineExceeded {
			metrics.Mark("worker.vm.provider.docker.boot.timeout")
		}
		return nil, ctx.Err()
	}

}

func (p *DockerProvider) imageForLanguage(language string) (string, error) {
	searchTag := "travis:" + language

	images, err := p.client.ListImages(docker.ListImagesOptions{All: true})
	if err != nil {
		return "", err
	}

	for _, image := range images {
		for _, tag := range image.RepoTags {
			if tag == searchTag {
				return image.ID, nil
			}
		}
	}

	return "", fmt.Errorf("no image found with language %s", language)
}

func (p *DockerProvider) checkoutCPUSets() (string, error) {
	p.cpuSetsMutex.Lock()
	defer p.cpuSetsMutex.Unlock()

	cpuSets := []int{}

	for i, checkedOut := range p.cpuSets {
		if !checkedOut {
			cpuSets = append(cpuSets, i)
		}
		if len(cpuSets) == 2 {
			break
		}
	}

	if len(cpuSets) != 2 {
		return "", fmt.Errorf("not enough free CPUsets")
	}

	p.cpuSets[cpuSets[0]] = true
	p.cpuSets[cpuSets[1]] = true

	return fmt.Sprintf("%d,%d", cpuSets[0], cpuSets[1]), nil
}

func (p *DockerProvider) checkinCPUSets(sets string) {
	p.cpuSetsMutex.Lock()
	defer p.cpuSetsMutex.Unlock()

	var cpu1, cpu2 int
	fmt.Sscanf(sets, "%d,%d", &cpu1, &cpu2)

	p.cpuSets[cpu1] = false
	p.cpuSets[cpu2] = false
}

func (i *DockerInstance) sshClient() (*ssh.Client, error) {
	var err error
	i.container, err = i.client.InspectContainer(i.container.ID)
	if err != nil {
		return nil, err
	}

	fmt.Printf("networksettings: %+v\n", i.container.NetworkSettings)

	time.Sleep(2 * time.Second)

	return ssh.Dial("tcp", fmt.Sprintf("%s:22", i.container.NetworkSettings.IPAddress), &ssh.ClientConfig{
		User: "travis",
		Auth: []ssh.AuthMethod{
			ssh.Password("travis"),
		},
	})
}

func (i *DockerInstance) UploadScript(ctx gocontext.Context, script []byte) error {
	client, err := i.sshClient()
	if err != nil {
		return err
	}
	defer client.Close()

	sftp, err := sftp.NewClient(client)
	if err != nil {
		return err
	}
	defer sftp.Close()

	f, err := sftp.Create("build.sh")
	if err != nil {
		return err
	}

	if _, err := f.Write(script); err != nil {
		return err
	}

	return nil
}

func (i *DockerInstance) RunScript(ctx gocontext.Context, output io.WriteCloser) (RunResult, error) {
	client, err := i.sshClient()
	if err != nil {
		return RunResult{Completed: false}, err
	}
	defer client.Close()

	session, err := client.NewSession()
	if err != nil {
		return RunResult{Completed: false}, err
	}
	defer session.Close()

	err = session.RequestPty("xterm", 80, 40, ssh.TerminalModes{})
	if err != nil {
		return RunResult{Completed: false}, err
	}

	session.Stdout = output
	session.Stderr = output

	err = session.Run("bash ~/build.sh")
	if err == nil {
		return RunResult{Completed: true, ExitCode: 0}, nil
	}

	switch err := err.(type) {
	case *ssh.ExitError:
		return RunResult{Completed: true, ExitCode: uint8(err.ExitStatus())}, nil
	default:
		return RunResult{Completed: false}, err
	}
}

func (i *DockerInstance) Stop(ctx gocontext.Context) error {
	defer i.provider.checkinCPUSets(i.container.Config.CPUSet)

	err := i.client.StopContainer(i.container.ID, 30)
	if err != nil {
		return err
	}

	return i.client.RemoveContainer(docker.RemoveContainerOptions{
		ID:            i.container.ID,
		RemoveVolumes: true,
		Force:         true,
	})
}
<commit_msg>Add support for TLS docker client<commit_after>package backend

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

	"code.google.com/p/go-uuid/uuid"

	"github.com/fsouza/go-dockerclient"
	"github.com/pkg/sftp"
	"github.com/travis-ci/worker/lib/context"
	"github.com/travis-ci/worker/lib/metrics"
	"golang.org/x/crypto/ssh"
	gocontext "golang.org/x/net/context"
)

var (
	errMissingEndpointConfig = fmt.Errorf("expected config key endpoint")
)

type DockerProvider struct {
	client *docker.Client

	cpuSetsMutex sync.Mutex
	cpuSets      []bool
}

type DockerInstance struct {
	client    *docker.Client
	provider  *DockerProvider
	container *docker.Container
}

func NewDockerProvider(config map[string]string) (*DockerProvider, error) {
	client, err := buildClient(config)
	if err != nil {
		return nil, err
	}

	cpuSetSize := runtime.NumCPU()
	if cpuSetSize < 2 {
		cpuSetSize = 2
	}

	return &DockerProvider{
		client:  client,
		cpuSets: make([]bool, cpuSetSize),
	}, nil
}

func buildClient(config map[string]string) (*docker.Client, error) {
	endpoint, ok := config["endpoint"]
	if !ok {
		return nil, errMissingEndpointConfig
	}

	if path, ok := config["cert_path"]; ok {
		ca := fmt.Sprintf("%s/ca.pem", path)
		cert := fmt.Sprintf("%s/cert.pem", path)
		key := fmt.Sprintf("%s/key.pem", path)
		return docker.NewTLSClient(endpoint, cert, key, ca)
	}

	return docker.NewClient(endpoint)
}

func (p *DockerProvider) Start(ctx gocontext.Context, startAttributes StartAttributes) (Instance, error) {
	cpuSets, err := p.checkoutCPUSets()
	if err != nil && cpuSets != "" {
		return nil, err
	}

	imageID, err := p.imageForLanguage(startAttributes.Language)
	if err != nil {
		return nil, err
	}

	dockerConfig := &docker.Config{
		Cmd:      []string{"/sbin/init"},
		Image:    imageID,
		Memory:   1024 * 1024 * 1024 * 4,
		Hostname: fmt.Sprintf("testing-go-%s", uuid.NewUUID()),
	}
	if cpuSets != "" {
		dockerConfig.CPUSet = cpuSets
	}

	container, err := p.client.CreateContainer(docker.CreateContainerOptions{Config: dockerConfig})
	if err != nil {
		if container != nil {
			err := p.client.RemoveContainer(docker.RemoveContainerOptions{
				ID:            container.ID,
				RemoveVolumes: true,
				Force:         true,
			})
			if err != nil {
				context.LoggerFromContext(ctx).WithField("err", err).Error("couldn't remove container after create failure")
			}
		}

		return nil, err
	}

	startBooting := time.Now()

	err = p.client.StartContainer(container.ID, &docker.HostConfig{})
	if err != nil {
		return nil, err
	}

	containerReady := make(chan *docker.Container)
	errChan := make(chan error)
	go func(id string) {
		for true {
			container, err := p.client.InspectContainer(id)
			if err != nil {
				errChan <- err
				return
			}

			if container.State.Running {
				containerReady <- container
				return
			}
		}
	}(container.ID)

	select {
	case container := <-containerReady:
		metrics.TimeSince("worker.vm.provider.docker.boot", startBooting)
		return &DockerInstance{
			client:    p.client,
			provider:  p,
			container: container,
		}, nil
	case err := <-errChan:
		return nil, err
	case <-ctx.Done():
		if ctx.Err() == gocontext.DeadlineExceeded {
			metrics.Mark("worker.vm.provider.docker.boot.timeout")
		}
		return nil, ctx.Err()
	}

}

func (p *DockerProvider) imageForLanguage(language string) (string, error) {
	searchTag := "travis:" + language

	images, err := p.client.ListImages(docker.ListImagesOptions{All: true})
	if err != nil {
		return "", err
	}

	for _, image := range images {
		for _, tag := range image.RepoTags {
			if tag == searchTag {
				return image.ID, nil
			}
		}
	}

	return "", fmt.Errorf("no image found with language %s", language)
}

func (p *DockerProvider) checkoutCPUSets() (string, error) {
	p.cpuSetsMutex.Lock()
	defer p.cpuSetsMutex.Unlock()

	cpuSets := []int{}

	for i, checkedOut := range p.cpuSets {
		if !checkedOut {
			cpuSets = append(cpuSets, i)
		}
		if len(cpuSets) == 2 {
			break
		}
	}

	if len(cpuSets) != 2 {
		return "", fmt.Errorf("not enough free CPUsets")
	}

	p.cpuSets[cpuSets[0]] = true
	p.cpuSets[cpuSets[1]] = true

	return fmt.Sprintf("%d,%d", cpuSets[0], cpuSets[1]), nil
}

func (p *DockerProvider) checkinCPUSets(sets string) {
	p.cpuSetsMutex.Lock()
	defer p.cpuSetsMutex.Unlock()

	var cpu1, cpu2 int
	fmt.Sscanf(sets, "%d,%d", &cpu1, &cpu2)

	p.cpuSets[cpu1] = false
	p.cpuSets[cpu2] = false
}

func (i *DockerInstance) sshClient() (*ssh.Client, error) {
	var err error
	i.container, err = i.client.InspectContainer(i.container.ID)
	if err != nil {
		return nil, err
	}

	fmt.Printf("networksettings: %+v\n", i.container.NetworkSettings)

	time.Sleep(2 * time.Second)

	return ssh.Dial("tcp", fmt.Sprintf("%s:22", i.container.NetworkSettings.IPAddress), &ssh.ClientConfig{
		User: "travis",
		Auth: []ssh.AuthMethod{
			ssh.Password("travis"),
		},
	})
}

func (i *DockerInstance) UploadScript(ctx gocontext.Context, script []byte) error {
	client, err := i.sshClient()
	if err != nil {
		return err
	}
	defer client.Close()

	sftp, err := sftp.NewClient(client)
	if err != nil {
		return err
	}
	defer sftp.Close()

	f, err := sftp.Create("build.sh")
	if err != nil {
		return err
	}

	if _, err := f.Write(script); err != nil {
		return err
	}

	return nil
}

func (i *DockerInstance) RunScript(ctx gocontext.Context, output io.WriteCloser) (RunResult, error) {
	client, err := i.sshClient()
	if err != nil {
		return RunResult{Completed: false}, err
	}
	defer client.Close()

	session, err := client.NewSession()
	if err != nil {
		return RunResult{Completed: false}, err
	}
	defer session.Close()

	err = session.RequestPty("xterm", 80, 40, ssh.TerminalModes{})
	if err != nil {
		return RunResult{Completed: false}, err
	}

	session.Stdout = output
	session.Stderr = output

	err = session.Run("bash ~/build.sh")
	if err == nil {
		return RunResult{Completed: true, ExitCode: 0}, nil
	}

	switch err := err.(type) {
	case *ssh.ExitError:
		return RunResult{Completed: true, ExitCode: uint8(err.ExitStatus())}, nil
	default:
		return RunResult{Completed: false}, err
	}
}

func (i *DockerInstance) Stop(ctx gocontext.Context) error {
	defer i.provider.checkinCPUSets(i.container.Config.CPUSet)

	err := i.client.StopContainer(i.container.ID, 30)
	if err != nil {
		return err
	}

	return i.client.RemoveContainer(docker.RemoveContainerOptions{
		ID:            i.container.ID,
		RemoveVolumes: true,
		Force:         true,
	})
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 OpenSDS 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 dorado

import (
	"strconv"
	"strings"

	log "github.com/golang/glog"
	. "github.com/opensds/opensds/contrib/drivers/utils/config"
	pb "github.com/opensds/opensds/pkg/dock/proto"
	"github.com/opensds/opensds/pkg/model"
	"github.com/opensds/opensds/pkg/utils/config"
)

const (
	defaultConfPath = "/etc/opensds/driver/huawei_dorado.yaml"
	defaultAZ       = "default"
	UnitGi          = 1024 * 1024 * 1024
)

type AuthOptions struct {
	Username  string `yaml:"userName,omitempty"`
	Password  string `yaml:"password,omitempty"`
	Endpoints string `yaml:"endpoints,omitempty"`
}

type DoradoConfig struct {
	AuthOptions `yaml:"authOptions"`
	Pool        map[string]PoolProperties `yaml:"pool,flow"`
}

type Driver struct {
	conf   *DoradoConfig
	client *DoradoClient
}

func (d *Driver) sector2Gb(sec string) int64 {
	capa, err := strconv.ParseInt(sec, 10, 64)
	if err != nil {
		log.Error("Convert capacity from string to number failed, error:", err)
		return 0
	}
	return capa * 512 / UnitGi
}

func (d *Driver) gb2Sector(gb int64) int64 {
	return gb * UnitGi / 512
}

func (d *Driver) Setup() error {
	// Read lvm config file
	conf := &DoradoConfig{}
	d.conf = conf
	path := config.CONF.OsdsDock.Backends.HuaweiDorado.ConfigPath

	if "" == path {
		path = defaultConfPath
	}
	Parse(conf, path)
	dp := strings.Split(conf.Endpoints, ",")
	client, err := NewClient(conf.Username, conf.Password, dp)
	d.client = client
	if err != nil {
		log.Errorf("Get new client failed, %v", err)
		return err
	}
	return nil
}

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

func (d *Driver) CreateVolume(opt *pb.CreateVolumeOpts) (*model.VolumeSpec, error) {
	//Convert the storage unit Giga to sector

	lun, err := d.client.CreateVolume(opt.GetName(), d.gb2Sector(opt.GetSize()), opt.GetDescription())
	if err != nil {
		log.Error("Create Volume Failed:", err)
		return nil, err
	}
	log.Infof("Create volume %s (%s) success.", opt.GetName(), lun.Id)
	return &model.VolumeSpec{
		BaseModel: &model.BaseModel{
			Id: lun.Id,
		},
		Name:             lun.Name,
		Size:             d.sector2Gb(lun.Capacity),
		Description:      lun.Description,
		AvailabilityZone: "dorado",
	}, nil
}

func (d *Driver) PullVolume(volID string) (*model.VolumeSpec, error) {
	lun, err := d.client.GetVolume(volID)
	if err != nil {
		return nil, err
	}

	return &model.VolumeSpec{
		BaseModel: &model.BaseModel{
			Id: lun.Id,
		},
		Name:             lun.Name,
		Size:             d.sector2Gb(lun.Capacity),
		Description:      lun.Description,
		AvailabilityZone: "dorado",
	}, nil
}

func (d *Driver) DeleteVolume(opt *pb.DeleteVolumeOpts) error {
	err := d.client.DeleteVolume(opt.Id)
	if err != nil {
		log.Errorf("Delete volume failed, volume id =%s , Error:%s", opt.GetId())
	}
	log.Info("Remove volume success, volume id =", opt.GetId())
	return nil
}

func (d *Driver) InitializeConnection(opt *pb.CreateAttachmentOpts) (*model.ConnectionInfo, error) {
	return &model.ConnectionInfo{}, nil
}

func (d *Driver) TerminateConnection(opt *pb.DeleteAttachmentOpts) error { return nil }

func (d *Driver) CreateSnapshot(opt *pb.CreateVolumeSnapshotOpts) (*model.VolumeSnapshotSpec, error) {
	snap, err := d.client.CreateSnapshot(opt.GetVolumeId(), opt.GetName(), opt.GetDescription())
	if err != nil {
		return nil, err
	}
	return &model.VolumeSnapshotSpec{
		BaseModel: &model.BaseModel{
			Id: snap.Id,
		},
		Name:        snap.Name,
		Description: snap.Description,
		VolumeId:    snap.ParentId,
		Size:        0,
	}, nil
}

func (d *Driver) PullSnapshot(id string) (*model.VolumeSnapshotSpec, error) {
	snap, err := d.client.GetSnapshot(id)
	if err != nil {
		return nil, err
	}
	return &model.VolumeSnapshotSpec{
		BaseModel: &model.BaseModel{
			Id: snap.Id,
		},
		Name:        snap.Name,
		Description: snap.Description,
		Size:        0,
		VolumeId:    snap.ParentId,
	}, nil
}

func (d *Driver) DeleteSnapshot(opt *pb.DeleteVolumeSnapshotOpts) error {
	err := d.client.DeleteSnapshot(opt.GetId())
	if err != nil {
		log.Errorf("Delete volume snapshot failed, volume id =%s , Error:%s", opt.GetId())
	}
	log.Info("Remove volume snapshot success, volume id =", opt.GetId())
	return nil
}

func (d *Driver) buildPoolParam(proper PoolProperties) map[string]interface{} {
	param := make(map[string]interface{})
	param["diskType"] = proper.DiskType
	param["iops"] = proper.IOPS
	param["bandwidth"] = proper.BandWidth
	return param
}

func (d *Driver) ListPools() ([]*model.StoragePoolSpec, error) {
	var pols []*model.StoragePoolSpec
	sp, err := d.client.ListStoragePools()
	if err != nil {
		return nil, err
	}
	for _, p := range sp {
		c := d.conf
		if _, ok := c.Pool[p.Name]; !ok {
			continue
		}
		param := d.buildPoolParam(c.Pool[p.Name])
		pol := &model.StoragePoolSpec{
			BaseModel: &model.BaseModel{
				Id: p.Id,
			},
			Name:             p.Name,
			TotalCapacity:    d.sector2Gb(p.UserTotalCapacity),
			FreeCapacity:     d.sector2Gb(p.UserFreeCapacity),
			Parameters:       param,
			AvailabilityZone: c.Pool[p.Name].AZ,
		}
		if pol.AvailabilityZone == "" {
			pol.AvailabilityZone = defaultAZ
		}
		pols = append(pols, pol)
	}
	return pols, nil
}

<commit_msg>Fix some bugs<commit_after>// Copyright (c) 2017 OpenSDS 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 dorado

import (
	"strconv"
	"strings"

	log "github.com/golang/glog"
	. "github.com/opensds/opensds/contrib/drivers/utils/config"
	pb "github.com/opensds/opensds/pkg/dock/proto"
	"github.com/opensds/opensds/pkg/model"
	"github.com/opensds/opensds/pkg/utils/config"
)

const (
	defaultConfPath = "/etc/opensds/driver/huawei_dorado.yaml"
	defaultAZ       = "default"
	UnitGi          = 1024 * 1024 * 1024
)

type AuthOptions struct {
	Username  string `yaml:"userName,omitempty"`
	Password  string `yaml:"password,omitempty"`
	Endpoints string `yaml:"endpoints,omitempty"`
}

type DoradoConfig struct {
	AuthOptions `yaml:"authOptions"`
	Pool        map[string]PoolProperties `yaml:"pool,flow"`
}

type Driver struct {
	conf   *DoradoConfig
	client *DoradoClient
}

func (d *Driver) sector2Gb(sec string) int64 {
	capa, err := strconv.ParseInt(sec, 10, 64)
	if err != nil {
		log.Error("Convert capacity from string to number failed, error:", err)
		return 0
	}
	return capa * 512 / UnitGi
}

func (d *Driver) gb2Sector(gb int64) int64 {
	return gb * UnitGi / 512
}

func (d *Driver) Setup() error {
	// Read huawei dorado config file
	conf := &DoradoConfig{}
	d.conf = conf
	path := config.CONF.OsdsDock.Backends.HuaweiDorado.ConfigPath

	if "" == path {
		path = defaultConfPath
	}
	Parse(conf, path)
	dp := strings.Split(conf.Endpoints, ",")
	client, err := NewClient(conf.Username, conf.Password, dp)
	d.client = client
	if err != nil {
		log.Errorf("Get new client failed, %v", err)
		return err
	}
	return nil
}

func (d *Driver) Unset() error {
	d.client.logout()
	return nil
}

func (d *Driver) CreateVolume(opt *pb.CreateVolumeOpts) (*model.VolumeSpec, error) {
	//Convert the storage unit Giga to sector

	lun, err := d.client.CreateVolume(opt.GetName(), d.gb2Sector(opt.GetSize()), opt.GetDescription())
	if err != nil {
		log.Error("Create Volume Failed:", err)
		return nil, err
	}
	log.Infof("Create volume %s (%s) success.", opt.GetName(), lun.Id)
	return &model.VolumeSpec{
		BaseModel: &model.BaseModel{
			Id: lun.Id,
		},
		Name:             lun.Name,
		Size:             d.sector2Gb(lun.Capacity),
		Description:      lun.Description,
		AvailabilityZone: "dorado",
	}, nil
}

func (d *Driver) PullVolume(volID string) (*model.VolumeSpec, error) {
	lun, err := d.client.GetVolume(volID)
	if err != nil {
		return nil, err
	}

	return &model.VolumeSpec{
		BaseModel: &model.BaseModel{
			Id: lun.Id,
		},
		Name:             lun.Name,
		Size:             d.sector2Gb(lun.Capacity),
		Description:      lun.Description,
		AvailabilityZone: "dorado",
	}, nil
}

func (d *Driver) DeleteVolume(opt *pb.DeleteVolumeOpts) error {
	err := d.client.DeleteVolume(opt.Id)
	if err != nil {
		log.Errorf("Delete volume failed, volume id =%s , Error:%s", opt.GetId())
	}
	log.Info("Remove volume success, volume id =", opt.GetId())
	return nil
}

func (d *Driver) InitializeConnection(opt *pb.CreateAttachmentOpts) (*model.ConnectionInfo, error) {
	return &model.ConnectionInfo{}, nil
}

func (d *Driver) TerminateConnection(opt *pb.DeleteAttachmentOpts) error { return nil }

func (d *Driver) CreateSnapshot(opt *pb.CreateVolumeSnapshotOpts) (*model.VolumeSnapshotSpec, error) {
	snap, err := d.client.CreateSnapshot(opt.GetVolumeId(), opt.GetName(), opt.GetDescription())
	if err != nil {
		return nil, err
	}
	return &model.VolumeSnapshotSpec{
		BaseModel: &model.BaseModel{
			Id: snap.Id,
		},
		Name:        snap.Name,
		Description: snap.Description,
		VolumeId:    snap.ParentId,
		Size:        0,
	}, nil
}

func (d *Driver) PullSnapshot(id string) (*model.VolumeSnapshotSpec, error) {
	snap, err := d.client.GetSnapshot(id)
	if err != nil {
		return nil, err
	}
	return &model.VolumeSnapshotSpec{
		BaseModel: &model.BaseModel{
			Id: snap.Id,
		},
		Name:        snap.Name,
		Description: snap.Description,
		Size:        0,
		VolumeId:    snap.ParentId,
	}, nil
}

func (d *Driver) DeleteSnapshot(opt *pb.DeleteVolumeSnapshotOpts) error {
	err := d.client.DeleteSnapshot(opt.GetId())
	if err != nil {
		log.Errorf("Delete volume snapshot failed, volume id =%s , Error:%s", opt.GetId())
	}
	log.Info("Remove volume snapshot success, volume id =", opt.GetId())
	return nil
}

func (d *Driver) buildPoolParam(proper PoolProperties) map[string]interface{} {
	param := make(map[string]interface{})
	param["diskType"] = proper.DiskType
	param["iops"] = proper.IOPS
	param["bandwidth"] = proper.BandWidth
	return param
}

func (d *Driver) ListPools() ([]*model.StoragePoolSpec, error) {
	var pols []*model.StoragePoolSpec
	sp, err := d.client.ListStoragePools()
	if err != nil {
		return nil, err
	}
	for _, p := range sp {
		c := d.conf
		if _, ok := c.Pool[p.Name]; !ok {
			continue
		}
		param := d.buildPoolParam(c.Pool[p.Name])
		pol := &model.StoragePoolSpec{
			BaseModel: &model.BaseModel{
				Id: p.Id,
			},
			Name:             p.Name,
			TotalCapacity:    d.sector2Gb(p.UserTotalCapacity),
			FreeCapacity:     d.sector2Gb(p.UserFreeCapacity),
			Parameters:       param,
			AvailabilityZone: c.Pool[p.Name].AZ,
		}
		if pol.AvailabilityZone == "" {
			pol.AvailabilityZone = defaultAZ
		}
		pols = append(pols, pol)
	}
	return pols, nil
}

<|endoftext|>


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

package main

import (
	"bufio"
	"exec"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"unicode"
	"utf8"
)

// Environment for commands.
var (
	XGC       []string // 6g -I _test -o _xtest_.6
	GC        []string // 6g -I _test _testmain.go
	GL        []string // 6l -L _test _testmain.6
	GOARCH    string
	GOROOT    string
	GORUN     string
	O         string
	args      []string // arguments passed to gotest; also passed to the binary
	fileNames []string
	env       = os.Environ()
)

// These strings are created by getTestNames.
var (
	insideFileNames  []string // list of *.go files inside the package.
	outsideFileNames []string // list of *.go files outside the package (in package foo_test).
)

var (
	files      []*File
	importPath string
)

// Flags for our own purposes. We do our own flag processing.
var (
	cFlag bool
	xFlag bool
)

// File represents a file that contains tests.
type File struct {
	name       string
	pkg        string
	file       *os.File
	astFile    *ast.File
	tests      []string // The names of the TestXXXs.
	benchmarks []string // The names of the BenchmarkXXXs.
}

func main() {
	flags()
	needMakefile()
	setEnvironment()
	getTestFileNames()
	parseFiles()
	getTestNames()
	run("gomake", "testpackage-clean")
	run("gomake", "testpackage", fmt.Sprintf("GOTESTFILES=%s", strings.Join(insideFileNames, " ")))
	if len(outsideFileNames) > 0 {
		run(append(XGC, outsideFileNames...)...)
	}
	importPath = runWithStdout("gomake", "-s", "importpath")
	writeTestmainGo()
	run(GC...)
	run(GL...)
	if !cFlag {
		runTestWithArgs("./" + O + ".out")
	}
}

// needMakefile tests that we have a Makefile in this directory.
func needMakefile() {
	if _, err := os.Stat("Makefile"); err != nil {
		Fatalf("please create a Makefile for gotest; see http://golang.org/doc/code.html for details")
	}
}

// Fatalf formats its arguments, prints the message with a final newline, and exits.
func Fatalf(s string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, "gotest: "+s+"\n", args...)
	os.Exit(2)
}

// theChar is the map from architecture to object character.
var theChar = map[string]string{
	"arm":   "5",
	"amd64": "6",
	"386":   "8",
}

// addEnv adds a name=value pair to the environment passed to subcommands.
// If the item is already in the environment, addEnv replaces the value.
func addEnv(name, value string) {
	for i := 0; i < len(env); i++ {
		if strings.HasPrefix(env[i], name+"=") {
			env[i] = name + "=" + value
			return
		}
	}
	env = append(env, name+"="+value)
}

// setEnvironment assembles the configuration for gotest and its subcommands.
func setEnvironment() {
	// Basic environment.
	GOROOT = runtime.GOROOT()
	addEnv("GOROOT", GOROOT)
	GOARCH = runtime.GOARCH
	addEnv("GOARCH", GOARCH)
	O = theChar[GOARCH]
	if O == "" {
		Fatalf("unknown architecture %s", GOARCH)
	}

	// Commands and their flags.
	gc := os.Getenv("GC")
	if gc == "" {
		gc = O + "g"
	}
	XGC = []string{gc, "-I", "_test", "-o", "_xtest_." + O}
	GC = []string{gc, "-I", "_test", "_testmain.go"}
	gl := os.Getenv("GL")
	if gl == "" {
		gl = O + "l"
	}
	GL = []string{gl, "-L", "_test", "_testmain." + O}

	// Silence make on Linux
	addEnv("MAKEFLAGS", "")
	addEnv("MAKELEVEL", "")
}

// getTestFileNames gets the set of files we're looking at.
// If gotest has no arguments, it scans the current directory for *_test.go files.
func getTestFileNames() {
	names := fileNames
	if len(names) == 0 {
		names, err = filepath.Glob("[^.]*_test.go")
		if err != nil {
			Fatalf("Glob pattern error: %s", err)
		}
		if len(names) == 0 {
			Fatalf(`no test files found: no match for "*_test.go"`)
		}
	}
	for _, n := range names {
		fd, err := os.Open(n, os.O_RDONLY, 0)
		if err != nil {
			Fatalf("%s: %s", n, err)
		}
		f := &File{name: n, file: fd}
		files = append(files, f)
	}
}

// parseFiles parses the files and remembers the packages we find.
func parseFiles() {
	fileSet := token.NewFileSet()
	for _, f := range files {
		// Report declaration errors so we can abort if the files are incorrect Go.
		file, err := parser.ParseFile(fileSet, f.name, nil, parser.DeclarationErrors)
		if err != nil {
			Fatalf("parse error: %s", err)
		}
		f.astFile = file
		f.pkg = file.Name.String()
		if f.pkg == "" {
			Fatalf("cannot happen: no package name in %s", f.name)
		}
	}
}

// getTestNames extracts the names of tests and benchmarks.  They are all
// top-level functions that are not methods.
func getTestNames() {
	for _, f := range files {
		for _, d := range f.astFile.Decls {
			n, ok := d.(*ast.FuncDecl)
			if !ok {
				continue
			}
			if n.Recv != nil { // a method, not a function.
				continue
			}
			name := n.Name.String()
			if isTest(name, "Test") {
				f.tests = append(f.tests, name)
			} else if isTest(name, "Benchmark") {
				f.benchmarks = append(f.benchmarks, name)
			}
			// TODO: worth checking the signature? Probably not.
		}
		if strings.HasSuffix(f.pkg, "_test") {
			outsideFileNames = append(outsideFileNames, f.name)
		} else {
			insideFileNames = append(insideFileNames, f.name)
		}
	}
}

// isTest tells whether name looks like a test (or benchmark, according to prefix).
// It is a Test (say) if there is a character after Test that is not a lower-case letter.
// We don't want TesticularCancer.
func isTest(name, prefix string) bool {
	if !strings.HasPrefix(name, prefix) {
		return false
	}
	if len(name) == len(prefix) { // "Test" is ok
		return true
	}
	rune, _ := utf8.DecodeRuneInString(name[len(prefix):])
	return !unicode.IsLower(rune)
}

func run(args ...string) {
	doRun(args, false)
}

// runWithStdout is like run, but returns the text of standard output with the last newline dropped.
func runWithStdout(argv ...string) string {
	s := doRun(argv, true)
	if len(s) == 0 {
		Fatalf("no output from command %s", strings.Join(argv, " "))
	}
	if s[len(s)-1] == '\n' {
		s = s[:len(s)-1]
	}
	return s
}

// runTestWithArgs appends gotest's runs the provided binary with the args passed on the command line.
func runTestWithArgs(binary string) {
	doRun(append([]string{binary}, args...), false)
}

// doRun is the general command runner.  The flag says whether we want to
// retrieve standard output.
func doRun(argv []string, returnStdout bool) string {
	if xFlag {
		fmt.Printf("gotest: %s\n", strings.Join(argv, " "))
	}
	if runtime.GOOS == "windows" && argv[0] == "gomake" {
		// gomake is a shell script and it cannot be executed directly on Windows.
		cmd := ""
		for i, v := range argv {
			if i > 0 {
				cmd += " "
			}
			cmd += `"` + v + `"`
		}
		argv = []string{"cmd", "/c", "sh", "-c", cmd}
	}
	var err os.Error
	argv[0], err = exec.LookPath(argv[0])
	if err != nil {
		Fatalf("can't find %s: %s", argv[0], err)
	}
	procAttr := &os.ProcAttr{
		Env: env,
		Files: []*os.File{
			os.Stdin,
			os.Stdout,
			os.Stderr,
		},
	}
	var r, w *os.File
	if returnStdout {
		r, w, err = os.Pipe()
		if err != nil {
			Fatalf("can't create pipe: %s", err)
		}
		procAttr.Files[1] = w
	}
	proc, err := os.StartProcess(argv[0], argv, procAttr)
	if err != nil {
		Fatalf("make failed to start: %s", err)
	}
	if returnStdout {
		defer r.Close()
		w.Close()
	}
	waitMsg, err := proc.Wait(0)
	if err != nil || waitMsg == nil {
		Fatalf("%s failed: %s", argv[0], err)
	}
	if !waitMsg.Exited() || waitMsg.ExitStatus() != 0 {
		Fatalf("%q failed: %s", strings.Join(argv, " "), waitMsg)
	}
	if returnStdout {
		b, err := ioutil.ReadAll(r)
		if err != nil {
			Fatalf("can't read output from command: %s", err)
		}
		return string(b)
	}
	return ""
}

// writeTestmainGo generates the test program to be compiled, "./_testmain.go".
func writeTestmainGo() {
	f, err := os.Open("_testmain.go", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0666)
	if err != nil {
		Fatalf("can't create _testmain.go: %s", err)
	}
	defer f.Close()
	b := bufio.NewWriter(f)
	defer b.Flush()

	// Package and imports.
	fmt.Fprint(b, "package main\n\n")
	// Are there tests from a package other than the one we're testing?
	// We can't just use file names because some of the things we compiled
	// contain no tests.
	outsideTests := false
	insideTests := false
	for _, f := range files {
		//println(f.name, f.pkg)
		if len(f.tests) == 0 && len(f.benchmarks) == 0 {
			continue
		}
		if strings.HasSuffix(f.pkg, "_test") {
			outsideTests = true
		} else {
			insideTests = true
		}
	}
	if insideTests {
		switch importPath {
		case "testing":
		case "main":
			// Import path main is reserved, so import with
			// explicit reference to ./_test/main instead.
			// Also, the file we are writing defines a function named main,
			// so rename this import to __main__ to avoid name conflict.
			fmt.Fprintf(b, "import __main__ %q\n", "./_test/main")
		default:
			fmt.Fprintf(b, "import %q\n", importPath)
		}
	}
	if outsideTests {
		fmt.Fprintf(b, "import %q\n", "./_xtest_")
	}
	fmt.Fprintf(b, "import %q\n", "testing")
	fmt.Fprintf(b, "import __os__     %q\n", "os")     // rename in case tested package is called os
	fmt.Fprintf(b, "import __regexp__ %q\n", "regexp") // rename in case tested package is called regexp
	fmt.Fprintln(b)                                    // for gofmt

	// Tests.
	fmt.Fprintln(b, "var tests = []testing.InternalTest{")
	for _, f := range files {
		for _, t := range f.tests {
			fmt.Fprintf(b, "\t{\"%s.%s\", %s.%s},\n", f.pkg, t, notMain(f.pkg), t)
		}
	}
	fmt.Fprintln(b, "}")
	fmt.Fprintln(b)

	// Benchmarks.
	fmt.Fprintln(b, "var benchmarks = []testing.InternalBenchmark{")
	for _, f := range files {
		for _, bm := range f.benchmarks {
			fmt.Fprintf(b, "\t{\"%s.%s\", %s.%s},\n", f.pkg, bm, notMain(f.pkg), bm)
		}
	}
	fmt.Fprintln(b, "}")

	// Body.
	fmt.Fprintln(b, testBody)
}

// notMain returns the package, renaming as appropriate if it's "main".
func notMain(pkg string) string {
	if pkg == "main" {
		return "__main__"
	}
	return pkg
}

// testBody is just copied to the output. It's the code that runs the tests.
var testBody = `
var matchPat string
var matchRe *__regexp__.Regexp

func matchString(pat, str string) (result bool, err __os__.Error) {
	if matchRe == nil || matchPat != pat {
		matchPat = pat
		matchRe, err = __regexp__.Compile(matchPat)
		if err != nil {
			return
		}
	}
	return matchRe.MatchString(str), nil
}

func main() {
	testing.Main(matchString, tests, benchmarks)
}`
<commit_msg>gotest: fix build<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bufio"
	"exec"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"unicode"
	"utf8"
)

// Environment for commands.
var (
	XGC       []string // 6g -I _test -o _xtest_.6
	GC        []string // 6g -I _test _testmain.go
	GL        []string // 6l -L _test _testmain.6
	GOARCH    string
	GOROOT    string
	GORUN     string
	O         string
	args      []string // arguments passed to gotest; also passed to the binary
	fileNames []string
	env       = os.Environ()
)

// These strings are created by getTestNames.
var (
	insideFileNames  []string // list of *.go files inside the package.
	outsideFileNames []string // list of *.go files outside the package (in package foo_test).
)

var (
	files      []*File
	importPath string
)

// Flags for our own purposes. We do our own flag processing.
var (
	cFlag bool
	xFlag bool
)

// File represents a file that contains tests.
type File struct {
	name       string
	pkg        string
	file       *os.File
	astFile    *ast.File
	tests      []string // The names of the TestXXXs.
	benchmarks []string // The names of the BenchmarkXXXs.
}

func main() {
	flags()
	needMakefile()
	setEnvironment()
	getTestFileNames()
	parseFiles()
	getTestNames()
	run("gomake", "testpackage-clean")
	run("gomake", "testpackage", fmt.Sprintf("GOTESTFILES=%s", strings.Join(insideFileNames, " ")))
	if len(outsideFileNames) > 0 {
		run(append(XGC, outsideFileNames...)...)
	}
	importPath = runWithStdout("gomake", "-s", "importpath")
	writeTestmainGo()
	run(GC...)
	run(GL...)
	if !cFlag {
		runTestWithArgs("./" + O + ".out")
	}
}

// needMakefile tests that we have a Makefile in this directory.
func needMakefile() {
	if _, err := os.Stat("Makefile"); err != nil {
		Fatalf("please create a Makefile for gotest; see http://golang.org/doc/code.html for details")
	}
}

// Fatalf formats its arguments, prints the message with a final newline, and exits.
func Fatalf(s string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, "gotest: "+s+"\n", args...)
	os.Exit(2)
}

// theChar is the map from architecture to object character.
var theChar = map[string]string{
	"arm":   "5",
	"amd64": "6",
	"386":   "8",
}

// addEnv adds a name=value pair to the environment passed to subcommands.
// If the item is already in the environment, addEnv replaces the value.
func addEnv(name, value string) {
	for i := 0; i < len(env); i++ {
		if strings.HasPrefix(env[i], name+"=") {
			env[i] = name + "=" + value
			return
		}
	}
	env = append(env, name+"="+value)
}

// setEnvironment assembles the configuration for gotest and its subcommands.
func setEnvironment() {
	// Basic environment.
	GOROOT = runtime.GOROOT()
	addEnv("GOROOT", GOROOT)
	GOARCH = runtime.GOARCH
	addEnv("GOARCH", GOARCH)
	O = theChar[GOARCH]
	if O == "" {
		Fatalf("unknown architecture %s", GOARCH)
	}

	// Commands and their flags.
	gc := os.Getenv("GC")
	if gc == "" {
		gc = O + "g"
	}
	XGC = []string{gc, "-I", "_test", "-o", "_xtest_." + O}
	GC = []string{gc, "-I", "_test", "_testmain.go"}
	gl := os.Getenv("GL")
	if gl == "" {
		gl = O + "l"
	}
	GL = []string{gl, "-L", "_test", "_testmain." + O}

	// Silence make on Linux
	addEnv("MAKEFLAGS", "")
	addEnv("MAKELEVEL", "")
}

// getTestFileNames gets the set of files we're looking at.
// If gotest has no arguments, it scans the current directory for *_test.go files.
func getTestFileNames() {
	names := fileNames
	if len(names) == 0 {
		var err os.Error
		names, err = filepath.Glob("[^.]*_test.go")
		if err != nil {
			Fatalf("Glob pattern error: %s", err)
		}
		if len(names) == 0 {
			Fatalf(`no test files found: no match for "*_test.go"`)
		}
	}
	for _, n := range names {
		fd, err := os.Open(n, os.O_RDONLY, 0)
		if err != nil {
			Fatalf("%s: %s", n, err)
		}
		f := &File{name: n, file: fd}
		files = append(files, f)
	}
}

// parseFiles parses the files and remembers the packages we find.
func parseFiles() {
	fileSet := token.NewFileSet()
	for _, f := range files {
		// Report declaration errors so we can abort if the files are incorrect Go.
		file, err := parser.ParseFile(fileSet, f.name, nil, parser.DeclarationErrors)
		if err != nil {
			Fatalf("parse error: %s", err)
		}
		f.astFile = file
		f.pkg = file.Name.String()
		if f.pkg == "" {
			Fatalf("cannot happen: no package name in %s", f.name)
		}
	}
}

// getTestNames extracts the names of tests and benchmarks.  They are all
// top-level functions that are not methods.
func getTestNames() {
	for _, f := range files {
		for _, d := range f.astFile.Decls {
			n, ok := d.(*ast.FuncDecl)
			if !ok {
				continue
			}
			if n.Recv != nil { // a method, not a function.
				continue
			}
			name := n.Name.String()
			if isTest(name, "Test") {
				f.tests = append(f.tests, name)
			} else if isTest(name, "Benchmark") {
				f.benchmarks = append(f.benchmarks, name)
			}
			// TODO: worth checking the signature? Probably not.
		}
		if strings.HasSuffix(f.pkg, "_test") {
			outsideFileNames = append(outsideFileNames, f.name)
		} else {
			insideFileNames = append(insideFileNames, f.name)
		}
	}
}

// isTest tells whether name looks like a test (or benchmark, according to prefix).
// It is a Test (say) if there is a character after Test that is not a lower-case letter.
// We don't want TesticularCancer.
func isTest(name, prefix string) bool {
	if !strings.HasPrefix(name, prefix) {
		return false
	}
	if len(name) == len(prefix) { // "Test" is ok
		return true
	}
	rune, _ := utf8.DecodeRuneInString(name[len(prefix):])
	return !unicode.IsLower(rune)
}

func run(args ...string) {
	doRun(args, false)
}

// runWithStdout is like run, but returns the text of standard output with the last newline dropped.
func runWithStdout(argv ...string) string {
	s := doRun(argv, true)
	if len(s) == 0 {
		Fatalf("no output from command %s", strings.Join(argv, " "))
	}
	if s[len(s)-1] == '\n' {
		s = s[:len(s)-1]
	}
	return s
}

// runTestWithArgs appends gotest's runs the provided binary with the args passed on the command line.
func runTestWithArgs(binary string) {
	doRun(append([]string{binary}, args...), false)
}

// doRun is the general command runner.  The flag says whether we want to
// retrieve standard output.
func doRun(argv []string, returnStdout bool) string {
	if xFlag {
		fmt.Printf("gotest: %s\n", strings.Join(argv, " "))
	}
	if runtime.GOOS == "windows" && argv[0] == "gomake" {
		// gomake is a shell script and it cannot be executed directly on Windows.
		cmd := ""
		for i, v := range argv {
			if i > 0 {
				cmd += " "
			}
			cmd += `"` + v + `"`
		}
		argv = []string{"cmd", "/c", "sh", "-c", cmd}
	}
	var err os.Error
	argv[0], err = exec.LookPath(argv[0])
	if err != nil {
		Fatalf("can't find %s: %s", argv[0], err)
	}
	procAttr := &os.ProcAttr{
		Env: env,
		Files: []*os.File{
			os.Stdin,
			os.Stdout,
			os.Stderr,
		},
	}
	var r, w *os.File
	if returnStdout {
		r, w, err = os.Pipe()
		if err != nil {
			Fatalf("can't create pipe: %s", err)
		}
		procAttr.Files[1] = w
	}
	proc, err := os.StartProcess(argv[0], argv, procAttr)
	if err != nil {
		Fatalf("make failed to start: %s", err)
	}
	if returnStdout {
		defer r.Close()
		w.Close()
	}
	waitMsg, err := proc.Wait(0)
	if err != nil || waitMsg == nil {
		Fatalf("%s failed: %s", argv[0], err)
	}
	if !waitMsg.Exited() || waitMsg.ExitStatus() != 0 {
		Fatalf("%q failed: %s", strings.Join(argv, " "), waitMsg)
	}
	if returnStdout {
		b, err := ioutil.ReadAll(r)
		if err != nil {
			Fatalf("can't read output from command: %s", err)
		}
		return string(b)
	}
	return ""
}

// writeTestmainGo generates the test program to be compiled, "./_testmain.go".
func writeTestmainGo() {
	f, err := os.Open("_testmain.go", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0666)
	if err != nil {
		Fatalf("can't create _testmain.go: %s", err)
	}
	defer f.Close()
	b := bufio.NewWriter(f)
	defer b.Flush()

	// Package and imports.
	fmt.Fprint(b, "package main\n\n")
	// Are there tests from a package other than the one we're testing?
	// We can't just use file names because some of the things we compiled
	// contain no tests.
	outsideTests := false
	insideTests := false
	for _, f := range files {
		//println(f.name, f.pkg)
		if len(f.tests) == 0 && len(f.benchmarks) == 0 {
			continue
		}
		if strings.HasSuffix(f.pkg, "_test") {
			outsideTests = true
		} else {
			insideTests = true
		}
	}
	if insideTests {
		switch importPath {
		case "testing":
		case "main":
			// Import path main is reserved, so import with
			// explicit reference to ./_test/main instead.
			// Also, the file we are writing defines a function named main,
			// so rename this import to __main__ to avoid name conflict.
			fmt.Fprintf(b, "import __main__ %q\n", "./_test/main")
		default:
			fmt.Fprintf(b, "import %q\n", importPath)
		}
	}
	if outsideTests {
		fmt.Fprintf(b, "import %q\n", "./_xtest_")
	}
	fmt.Fprintf(b, "import %q\n", "testing")
	fmt.Fprintf(b, "import __os__     %q\n", "os")     // rename in case tested package is called os
	fmt.Fprintf(b, "import __regexp__ %q\n", "regexp") // rename in case tested package is called regexp
	fmt.Fprintln(b)                                    // for gofmt

	// Tests.
	fmt.Fprintln(b, "var tests = []testing.InternalTest{")
	for _, f := range files {
		for _, t := range f.tests {
			fmt.Fprintf(b, "\t{\"%s.%s\", %s.%s},\n", f.pkg, t, notMain(f.pkg), t)
		}
	}
	fmt.Fprintln(b, "}")
	fmt.Fprintln(b)

	// Benchmarks.
	fmt.Fprintln(b, "var benchmarks = []testing.InternalBenchmark{")
	for _, f := range files {
		for _, bm := range f.benchmarks {
			fmt.Fprintf(b, "\t{\"%s.%s\", %s.%s},\n", f.pkg, bm, notMain(f.pkg), bm)
		}
	}
	fmt.Fprintln(b, "}")

	// Body.
	fmt.Fprintln(b, testBody)
}

// notMain returns the package, renaming as appropriate if it's "main".
func notMain(pkg string) string {
	if pkg == "main" {
		return "__main__"
	}
	return pkg
}

// testBody is just copied to the output. It's the code that runs the tests.
var testBody = `
var matchPat string
var matchRe *__regexp__.Regexp

func matchString(pat, str string) (result bool, err __os__.Error) {
	if matchRe == nil || matchPat != pat {
		matchPat = pat
		matchRe, err = __regexp__.Compile(matchPat)
		if err != nil {
			return
		}
	}
	return matchRe.MatchString(str), nil
}

func main() {
	testing.Main(matchString, tests, benchmarks)
}`
<|endoftext|>


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

package packagestest

import (
	"fmt"
	"go/token"
	"path/filepath"
	"reflect"
	"regexp"

	"golang.org/x/tools/go/expect"
	"golang.org/x/tools/go/packages"
)

const (
	markMethod    = "mark"
	eofIdentifier = "EOF"
)

// Expect invokes the supplied methods for all expectation notes found in
// the exported source files.
//
// All exported go source files are parsed to collect the expectation
// notes.
// See the documentation for expect.Parse for how the notes are collected
// and parsed.
//
// The methods are supplied as a map of name to function, and those functions
// will be matched against the expectations by name.
// Notes with no matching function will be skipped, and functions with no
// matching notes will not be invoked.
// If there are no registered markers yet, a special pass will be run first
// which adds any markers declared with @mark(Name, pattern) or @name. These
// call the Mark method to add the marker to the global set.
// You can register the "mark" method to override these in your own call to
// Expect. The bound Mark function is usable directly in your method map, so
//    exported.Expect(map[string]interface{}{"mark": exported.Mark})
// replicates the built in behavior.
//
// Method invocation
//
// When invoking a method the expressions in the parameter list need to be
// converted to values to be passed to the method.
// There are a very limited set of types the arguments are allowed to be.
//   expect.Comment : passed the Comment instance being evaluated.
//   string : can be supplied either a string literal or an identifier.
//   int : can only be supplied an integer literal.
//   token.Pos : has a file position calculated as described below.
//   token.Position : has a file position calculated as described below.
//
// Position calculation
//
// There is some extra handling when a parameter is being coerced into a
// token.Pos, token.Position or Range type argument.
//
// If the parameter is an identifier, it will be treated as the name of an
// marker to look up (as if markers were global variables).
//
// If it is a string or regular expression, then it will be passed to
// expect.MatchBefore to look up a match in the line at which it was declared.
//
// It is safe to call this repeatedly with different method sets, but it is
// not safe to call it concurrently.
func (e *Exported) Expect(methods map[string]interface{}) error {
	if err := e.getNotes(); err != nil {
		return err
	}
	if err := e.getMarkers(); err != nil {
		return err
	}
	var err error
	ms := make(map[string]method, len(methods))
	for name, f := range methods {
		mi := method{f: reflect.ValueOf(f)}
		mi.converters = make([]converter, mi.f.Type().NumIn())
		for i := 0; i < len(mi.converters); i++ {
			mi.converters[i], err = e.buildConverter(mi.f.Type().In(i))
			if err != nil {
				return fmt.Errorf("invalid method %v: %v", name, err)
			}
		}
		ms[name] = mi
	}
	for _, n := range e.notes {
		if n.Args == nil {
			// simple identifier form, convert to a call to mark
			n = &expect.Note{
				Pos:  n.Pos,
				Name: markMethod,
				Args: []interface{}{n.Name, n.Name},
			}
		}
		mi, ok := ms[n.Name]
		if !ok {
			continue
		}
		params := make([]reflect.Value, len(mi.converters))
		args := n.Args
		for i, convert := range mi.converters {
			params[i], args, err = convert(n, args)
			if err != nil {
				return fmt.Errorf("%v: %v", e.fset.Position(n.Pos), err)
			}
		}
		if len(args) > 0 {
			return fmt.Errorf("%v: unwanted args got %+v extra", e.fset.Position(n.Pos), args)
		}
		//TODO: catch the error returned from the method
		mi.f.Call(params)
	}
	return nil
}

type Range struct {
	Start token.Pos
	End   token.Pos
}

// Mark adds a new marker to the known set.
func (e *Exported) Mark(name string, r Range) {
	if e.markers == nil {
		e.markers = make(map[string]Range)
	}
	e.markers[name] = r
}

func (e *Exported) getNotes() error {
	if e.notes != nil {
		return nil
	}
	notes := []*expect.Note{}
	var dirs []string
	for _, module := range e.written {
		for _, filename := range module {
			dirs = append(dirs, filepath.Dir(filename))
		}
	}
	pkgs, err := packages.Load(e.Config, dirs...)
	if err != nil {
		return fmt.Errorf("unable to load packages for directories %s: %v", dirs, err)
	}
	for _, pkg := range pkgs {
		for _, filename := range pkg.GoFiles {
			l, err := expect.Parse(e.fset, filename, nil)
			if err != nil {
				return fmt.Errorf("Failed to extract expectations: %v", err)
			}
			notes = append(notes, l...)
		}
	}
	e.notes = notes
	return nil
}

func (e *Exported) getMarkers() error {
	if e.markers != nil {
		return nil
	}
	// set markers early so that we don't call getMarkers again from Expect
	e.markers = make(map[string]Range)
	return e.Expect(map[string]interface{}{
		markMethod: e.Mark,
	})
}

var (
	noteType       = reflect.TypeOf((*expect.Note)(nil))
	identifierType = reflect.TypeOf(expect.Identifier(""))
	posType        = reflect.TypeOf(token.Pos(0))
	positionType   = reflect.TypeOf(token.Position{})
	rangeType      = reflect.TypeOf(Range{})
	fsetType       = reflect.TypeOf((*token.FileSet)(nil))
)

// converter converts from a marker's argument parsed from the comment to
// reflect values passed to the method during Invoke.
// It takes the args remaining, and returns the args it did not consume.
// This allows a converter to consume 0 args for well known types, or multiple
// args for compound types.
type converter func(*expect.Note, []interface{}) (reflect.Value, []interface{}, error)

// method is used to track information about Invoke methods that is expensive to
// calculate so that we can work it out once rather than per marker.
type method struct {
	f          reflect.Value // the reflect value of the passed in method
	converters []converter   // the parameter converters for the method
}

// buildConverter works out what function should be used to go from an ast expressions to a reflect
// value of the type expected by a method.
// It is called when only the target type is know, it returns converters that are flexible across
// all supported expression types for that target type.
func (e *Exported) buildConverter(pt reflect.Type) (converter, error) {
	switch {
	case pt == noteType:
		return func(n *expect.Note, args []interface{}) (reflect.Value, []interface{}, error) {
			return reflect.ValueOf(n), args, nil
		}, nil
	case pt == fsetType:
		return func(n *expect.Note, args []interface{}) (reflect.Value, []interface{}, error) {
			return reflect.ValueOf(e.fset), args, nil
		}, nil
	case pt == posType:
		return func(n *expect.Note, args []interface{}) (reflect.Value, []interface{}, error) {
			r, remains, err := e.rangeConverter(n, args)
			if err != nil {
				return reflect.Value{}, nil, err
			}
			return reflect.ValueOf(r.Start), remains, nil
		}, nil
	case pt == positionType:
		return func(n *expect.Note, args []interface{}) (reflect.Value, []interface{}, error) {
			r, remains, err := e.rangeConverter(n, args)
			if err != nil {
				return reflect.Value{}, nil, err
			}
			return reflect.ValueOf(e.fset.Position(r.Start)), remains, nil
		}, nil
	case pt == rangeType:
		return func(n *expect.Note, args []interface{}) (reflect.Value, []interface{}, error) {
			r, remains, err := e.rangeConverter(n, args)
			if err != nil {
				return reflect.Value{}, nil, err
			}
			return reflect.ValueOf(r), remains, nil
		}, nil
	case pt == identifierType:
		return func(n *expect.Note, args []interface{}) (reflect.Value, []interface{}, error) {
			arg := args[0]
			args = args[1:]
			switch arg := arg.(type) {
			case expect.Identifier:
				return reflect.ValueOf(arg), args, nil
			default:
				return reflect.Value{}, nil, fmt.Errorf("cannot convert %v to string", arg)
			}
		}, nil
	case pt.Kind() == reflect.String:
		return func(n *expect.Note, args []interface{}) (reflect.Value, []interface{}, error) {
			arg := args[0]
			args = args[1:]
			switch arg := arg.(type) {
			case expect.Identifier:
				return reflect.ValueOf(string(arg)), args, nil
			case string:
				return reflect.ValueOf(arg), args, nil
			default:
				return reflect.Value{}, nil, fmt.Errorf("cannot convert %v to string", arg)
			}
		}, nil
	case pt.Kind() == reflect.Int64:
		return func(n *expect.Note, args []interface{}) (reflect.Value, []interface{}, error) {
			arg := args[0]
			args = args[1:]
			switch arg := arg.(type) {
			case int64:
				return reflect.ValueOf(arg), args, nil
			default:
				return reflect.Value{}, nil, fmt.Errorf("cannot convert %v to int", arg)
			}
		}, nil
	case pt.Kind() == reflect.Bool:
		return func(n *expect.Note, args []interface{}) (reflect.Value, []interface{}, error) {
			arg := args[0]
			args = args[1:]
			b, ok := arg.(bool)
			if !ok {
				return reflect.Value{}, nil, fmt.Errorf("cannot convert %v to bool", arg)
			}
			return reflect.ValueOf(b), args, nil
		}, nil
	case pt.Kind() == reflect.Slice:
		return func(n *expect.Note, args []interface{}) (reflect.Value, []interface{}, error) {
			converter, err := e.buildConverter(pt.Elem())
			if err != nil {
				return reflect.Value{}, nil, err
			}
			result := reflect.MakeSlice(reflect.SliceOf(pt.Elem()), 0, len(args))
			for range args {
				value, remains, err := converter(n, args)
				if err != nil {
					return reflect.Value{}, nil, err
				}
				result = reflect.Append(result, value)
				args = remains
			}
			return result, args, nil
		}, nil
	default:
		return nil, fmt.Errorf("param has invalid type %v", pt)
	}
}

func (e *Exported) rangeConverter(n *expect.Note, args []interface{}) (Range, []interface{}, error) {
	if len(args) < 1 {
		return Range{}, nil, fmt.Errorf("missing argument")
	}
	arg := args[0]
	args = args[1:]
	switch arg := arg.(type) {
	case expect.Identifier:
		// handle the special identifiers
		switch arg {
		case eofIdentifier:
			// end of file identifier, look up the current file
			f := e.fset.File(n.Pos)
			eof := f.Pos(f.Size())
			return Range{Start: eof, End: token.NoPos}, args, nil
		default:
			// look up an marker by name
			mark, ok := e.markers[string(arg)]
			if !ok {
				return Range{}, nil, fmt.Errorf("cannot find marker %v", arg)
			}
			return mark, args, nil
		}
	case string:
		start, end, err := expect.MatchBefore(e.fset, e.fileContents, n.Pos, arg)
		if err != nil {
			return Range{}, nil, err
		}
		if start == token.NoPos {
			return Range{}, nil, fmt.Errorf("%v: pattern %s did not match", e.fset.Position(n.Pos), arg)
		}
		return Range{Start: start, End: end}, args, nil
	case *regexp.Regexp:
		start, end, err := expect.MatchBefore(e.fset, e.fileContents, n.Pos, arg)
		if err != nil {
			return Range{}, nil, err
		}
		if start == token.NoPos {
			return Range{}, nil, fmt.Errorf("%v: pattern %s did not match", e.fset.Position(n.Pos), arg)
		}
		return Range{Start: start, End: end}, args, nil
	default:
		return Range{}, nil, fmt.Errorf("cannot convert %v to pos", arg)
	}
}
<commit_msg>go/packages/packagestest: allow expectations to have *regexp.Regexp and interface{} parameters<commit_after>// Copyright 2018 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package packagestest

import (
	"fmt"
	"go/token"
	"path/filepath"
	"reflect"
	"regexp"

	"golang.org/x/tools/go/expect"
	"golang.org/x/tools/go/packages"
)

const (
	markMethod    = "mark"
	eofIdentifier = "EOF"
)

// Expect invokes the supplied methods for all expectation notes found in
// the exported source files.
//
// All exported go source files are parsed to collect the expectation
// notes.
// See the documentation for expect.Parse for how the notes are collected
// and parsed.
//
// The methods are supplied as a map of name to function, and those functions
// will be matched against the expectations by name.
// Notes with no matching function will be skipped, and functions with no
// matching notes will not be invoked.
// If there are no registered markers yet, a special pass will be run first
// which adds any markers declared with @mark(Name, pattern) or @name. These
// call the Mark method to add the marker to the global set.
// You can register the "mark" method to override these in your own call to
// Expect. The bound Mark function is usable directly in your method map, so
//    exported.Expect(map[string]interface{}{"mark": exported.Mark})
// replicates the built in behavior.
//
// Method invocation
//
// When invoking a method the expressions in the parameter list need to be
// converted to values to be passed to the method.
// There are a very limited set of types the arguments are allowed to be.
//   expect.Comment : passed the Comment instance being evaluated.
//   string : can be supplied either a string literal or an identifier.
//   int : can only be supplied an integer literal.
//   *regexp.Regexp : can only be supplied a regular expression literal
//   token.Pos : has a file position calculated as described below.
//   token.Position : has a file position calculated as described below.
//   interface{} : will be passed any value
//
// Position calculation
//
// There is some extra handling when a parameter is being coerced into a
// token.Pos, token.Position or Range type argument.
//
// If the parameter is an identifier, it will be treated as the name of an
// marker to look up (as if markers were global variables).
//
// If it is a string or regular expression, then it will be passed to
// expect.MatchBefore to look up a match in the line at which it was declared.
//
// It is safe to call this repeatedly with different method sets, but it is
// not safe to call it concurrently.
func (e *Exported) Expect(methods map[string]interface{}) error {
	if err := e.getNotes(); err != nil {
		return err
	}
	if err := e.getMarkers(); err != nil {
		return err
	}
	var err error
	ms := make(map[string]method, len(methods))
	for name, f := range methods {
		mi := method{f: reflect.ValueOf(f)}
		mi.converters = make([]converter, mi.f.Type().NumIn())
		for i := 0; i < len(mi.converters); i++ {
			mi.converters[i], err = e.buildConverter(mi.f.Type().In(i))
			if err != nil {
				return fmt.Errorf("invalid method %v: %v", name, err)
			}
		}
		ms[name] = mi
	}
	for _, n := range e.notes {
		if n.Args == nil {
			// simple identifier form, convert to a call to mark
			n = &expect.Note{
				Pos:  n.Pos,
				Name: markMethod,
				Args: []interface{}{n.Name, n.Name},
			}
		}
		mi, ok := ms[n.Name]
		if !ok {
			continue
		}
		params := make([]reflect.Value, len(mi.converters))
		args := n.Args
		for i, convert := range mi.converters {
			params[i], args, err = convert(n, args)
			if err != nil {
				return fmt.Errorf("%v: %v", e.fset.Position(n.Pos), err)
			}
		}
		if len(args) > 0 {
			return fmt.Errorf("%v: unwanted args got %+v extra", e.fset.Position(n.Pos), args)
		}
		//TODO: catch the error returned from the method
		mi.f.Call(params)
	}
	return nil
}

type Range struct {
	Start token.Pos
	End   token.Pos
}

// Mark adds a new marker to the known set.
func (e *Exported) Mark(name string, r Range) {
	if e.markers == nil {
		e.markers = make(map[string]Range)
	}
	e.markers[name] = r
}

func (e *Exported) getNotes() error {
	if e.notes != nil {
		return nil
	}
	notes := []*expect.Note{}
	var dirs []string
	for _, module := range e.written {
		for _, filename := range module {
			dirs = append(dirs, filepath.Dir(filename))
		}
	}
	pkgs, err := packages.Load(e.Config, dirs...)
	if err != nil {
		return fmt.Errorf("unable to load packages for directories %s: %v", dirs, err)
	}
	for _, pkg := range pkgs {
		for _, filename := range pkg.GoFiles {
			l, err := expect.Parse(e.fset, filename, nil)
			if err != nil {
				return fmt.Errorf("Failed to extract expectations: %v", err)
			}
			notes = append(notes, l...)
		}
	}
	e.notes = notes
	return nil
}

func (e *Exported) getMarkers() error {
	if e.markers != nil {
		return nil
	}
	// set markers early so that we don't call getMarkers again from Expect
	e.markers = make(map[string]Range)
	return e.Expect(map[string]interface{}{
		markMethod: e.Mark,
	})
}

var (
	noteType       = reflect.TypeOf((*expect.Note)(nil))
	identifierType = reflect.TypeOf(expect.Identifier(""))
	posType        = reflect.TypeOf(token.Pos(0))
	positionType   = reflect.TypeOf(token.Position{})
	rangeType      = reflect.TypeOf(Range{})
	fsetType       = reflect.TypeOf((*token.FileSet)(nil))
	regexType      = reflect.TypeOf((*regexp.Regexp)(nil))
)

// converter converts from a marker's argument parsed from the comment to
// reflect values passed to the method during Invoke.
// It takes the args remaining, and returns the args it did not consume.
// This allows a converter to consume 0 args for well known types, or multiple
// args for compound types.
type converter func(*expect.Note, []interface{}) (reflect.Value, []interface{}, error)

// method is used to track information about Invoke methods that is expensive to
// calculate so that we can work it out once rather than per marker.
type method struct {
	f          reflect.Value // the reflect value of the passed in method
	converters []converter   // the parameter converters for the method
}

// buildConverter works out what function should be used to go from an ast expressions to a reflect
// value of the type expected by a method.
// It is called when only the target type is know, it returns converters that are flexible across
// all supported expression types for that target type.
func (e *Exported) buildConverter(pt reflect.Type) (converter, error) {
	switch {
	case pt == noteType:
		return func(n *expect.Note, args []interface{}) (reflect.Value, []interface{}, error) {
			return reflect.ValueOf(n), args, nil
		}, nil
	case pt == fsetType:
		return func(n *expect.Note, args []interface{}) (reflect.Value, []interface{}, error) {
			return reflect.ValueOf(e.fset), args, nil
		}, nil
	case pt == posType:
		return func(n *expect.Note, args []interface{}) (reflect.Value, []interface{}, error) {
			r, remains, err := e.rangeConverter(n, args)
			if err != nil {
				return reflect.Value{}, nil, err
			}
			return reflect.ValueOf(r.Start), remains, nil
		}, nil
	case pt == positionType:
		return func(n *expect.Note, args []interface{}) (reflect.Value, []interface{}, error) {
			r, remains, err := e.rangeConverter(n, args)
			if err != nil {
				return reflect.Value{}, nil, err
			}
			return reflect.ValueOf(e.fset.Position(r.Start)), remains, nil
		}, nil
	case pt == rangeType:
		return func(n *expect.Note, args []interface{}) (reflect.Value, []interface{}, error) {
			r, remains, err := e.rangeConverter(n, args)
			if err != nil {
				return reflect.Value{}, nil, err
			}
			return reflect.ValueOf(r), remains, nil
		}, nil
	case pt == identifierType:
		return func(n *expect.Note, args []interface{}) (reflect.Value, []interface{}, error) {
			arg := args[0]
			args = args[1:]
			switch arg := arg.(type) {
			case expect.Identifier:
				return reflect.ValueOf(arg), args, nil
			default:
				return reflect.Value{}, nil, fmt.Errorf("cannot convert %v to string", arg)
			}
		}, nil

	case pt == regexType:
		return func(n *expect.Note, args []interface{}) (reflect.Value, []interface{}, error) {
			arg := args[0]
			args = args[1:]
			if _, ok := arg.(*regexp.Regexp); !ok {
				return reflect.Value{}, nil, fmt.Errorf("cannot convert %v to *regexp.Regexp", arg)
			}
			return reflect.ValueOf(arg), args, nil
		}, nil

	case pt.Kind() == reflect.String:
		return func(n *expect.Note, args []interface{}) (reflect.Value, []interface{}, error) {
			arg := args[0]
			args = args[1:]
			switch arg := arg.(type) {
			case expect.Identifier:
				return reflect.ValueOf(string(arg)), args, nil
			case string:
				return reflect.ValueOf(arg), args, nil
			default:
				return reflect.Value{}, nil, fmt.Errorf("cannot convert %v to string", arg)
			}
		}, nil
	case pt.Kind() == reflect.Int64:
		return func(n *expect.Note, args []interface{}) (reflect.Value, []interface{}, error) {
			arg := args[0]
			args = args[1:]
			switch arg := arg.(type) {
			case int64:
				return reflect.ValueOf(arg), args, nil
			default:
				return reflect.Value{}, nil, fmt.Errorf("cannot convert %v to int", arg)
			}
		}, nil
	case pt.Kind() == reflect.Bool:
		return func(n *expect.Note, args []interface{}) (reflect.Value, []interface{}, error) {
			arg := args[0]
			args = args[1:]
			b, ok := arg.(bool)
			if !ok {
				return reflect.Value{}, nil, fmt.Errorf("cannot convert %v to bool", arg)
			}
			return reflect.ValueOf(b), args, nil
		}, nil
	case pt.Kind() == reflect.Slice:
		return func(n *expect.Note, args []interface{}) (reflect.Value, []interface{}, error) {
			converter, err := e.buildConverter(pt.Elem())
			if err != nil {
				return reflect.Value{}, nil, err
			}
			result := reflect.MakeSlice(reflect.SliceOf(pt.Elem()), 0, len(args))
			for range args {
				value, remains, err := converter(n, args)
				if err != nil {
					return reflect.Value{}, nil, err
				}
				result = reflect.Append(result, value)
				args = remains
			}
			return result, args, nil
		}, nil
	default:
		if pt.Kind() == reflect.Interface && pt.NumMethod() == 0 {
			return func(n *expect.Note, args []interface{}) (reflect.Value, []interface{}, error) {
				return reflect.ValueOf(args[0]), args[1:], nil
			}, nil
		}
		return nil, fmt.Errorf("param has unexpected type %v (kind %v)", pt, pt.Kind())
	}
}

func (e *Exported) rangeConverter(n *expect.Note, args []interface{}) (Range, []interface{}, error) {
	if len(args) < 1 {
		return Range{}, nil, fmt.Errorf("missing argument")
	}
	arg := args[0]
	args = args[1:]
	switch arg := arg.(type) {
	case expect.Identifier:
		// handle the special identifiers
		switch arg {
		case eofIdentifier:
			// end of file identifier, look up the current file
			f := e.fset.File(n.Pos)
			eof := f.Pos(f.Size())
			return Range{Start: eof, End: token.NoPos}, args, nil
		default:
			// look up an marker by name
			mark, ok := e.markers[string(arg)]
			if !ok {
				return Range{}, nil, fmt.Errorf("cannot find marker %v", arg)
			}
			return mark, args, nil
		}
	case string:
		start, end, err := expect.MatchBefore(e.fset, e.fileContents, n.Pos, arg)
		if err != nil {
			return Range{}, nil, err
		}
		if start == token.NoPos {
			return Range{}, nil, fmt.Errorf("%v: pattern %s did not match", e.fset.Position(n.Pos), arg)
		}
		return Range{Start: start, End: end}, args, nil
	case *regexp.Regexp:
		start, end, err := expect.MatchBefore(e.fset, e.fileContents, n.Pos, arg)
		if err != nil {
			return Range{}, nil, err
		}
		if start == token.NoPos {
			return Range{}, nil, fmt.Errorf("%v: pattern %s did not match", e.fset.Position(n.Pos), arg)
		}
		return Range{Start: start, End: end}, args, nil
	default:
		return Range{}, nil, fmt.Errorf("cannot convert %v to pos", arg)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/hermantai/samples/go/commonutil"
	"golang.org/x/tour/tree"
	"sync"
	"time"
)

func main() {
	commonutil.PrintSection("Goroutines")
	goroutinesSample()

	commonutil.PrintSection("Channels")
	channelsSample()

	commonutil.PrintSection("Equivalent binary tree")
	fmt.Println(Same(tree.New(1), tree.New(1)))
	fmt.Println(Same(tree.New(1), tree.New(2)))

	commonutil.PrintSection("Locks")
	locksSample()
}

// end of main

func goroutinesSample() {
	go say("world")
	say("hello")
}

func say(s string) {
	for i := 0; i < 5; i++ {
		time.Sleep(100 * time.Millisecond)
		fmt.Println(s)
	}
}

func channelsSample() {
	s := []int{7, 2, 8, -9, 4, 0}

	c := make(chan int)
	go sum(s[:len(s)/2], c)
	go sum(s[len(s)/2:], c)

	x, y := <-c, <-c

	fmt.Println(x, y, x+y)

	commonutil.PrintSubsection("Buffered channels")
	ch := make(chan int, 2)
	ch <- 3
	ch <- 4
	fmt.Println(<-ch, <-ch)

	ch2 := make(chan int, 10)
	go fibonacci(cap(ch2), ch2)
	for i := range ch2 {
		fmt.Println(i)
	}

	commonutil.PrintSubsection("Select")
	ch3 := make(chan int)
	quit := make(chan int)
	go func() {
		for i := 0; i < 10; i++ {
			fmt.Println(<-ch3)
		}
		quit <- 0
	}()
	fibWithSelect(ch3, quit)

	commonutil.PrintSubsection("default selection")
	tick := time.Tick(100 * time.Millisecond)
	boom := time.After(500 * time.Millisecond)
	for {
		select {
		case <-tick:
			fmt.Println("tick.")
		case <-boom:
			fmt.Println("BOOM!")
			return
		default:
			fmt.Println("    .")
			time.Sleep(50 * time.Millisecond)
		}
	}
}

func sum(s []int, c chan int) {
	sum := 0
	for _, v := range s {
		sum += v
	}
	c <- sum
}

func fibonacci(n int, c chan int) {
	x, y := 0, 1
	for i := 0; i < n; i++ {
		c <- x
		x, y = y, x+y
	}
	close(c)
}

func fibWithSelect(c, quit chan int) {
	x, y := 0, 1
	for {
		select {
		case c <- x:
			x, y = y, x+y
		case <-quit:
			fmt.Println("quit")
			return
		}
	}
}

// Walk walks the tree t sending all values
// from the tree to the channel ch.
func Walk(t *tree.Tree, ch chan int) {
	if t == nil {
		return
	}
	Walk(t.Left, ch)
	ch <- t.Value
	Walk(t.Right, ch)
}

// Same determines whether the trees
// t1 and t2 contain the same values.
func Same(t1, t2 *tree.Tree) bool {
	ch1, ch2 := make(chan int), make(chan int)
	go func() {
		Walk(t1, ch1)
		close(ch1)
	}()
	go func() {
		Walk(t2, ch2)
		close(ch2)
	}()

	for {
		x, ok := <-ch1
		if !ok {
			_, ok := <-ch2
			return !ok
		}

		if y, ok := <-ch2; !ok || x != y {
			return false
		}
	}
}

func locksSample() {
	c := SafeCounter{v: make(map[string]int)}
	for i := 0; i < 1000; i++ {
		go c.Inc("somekey")
	}

	time.Sleep(time.Second)
	fmt.Println(c.Value("somekey"))
}

// SafeCounter is safe to use concurrently.
type SafeCounter struct {
	v   map[string]int
	mux sync.Mutex
}

// Inc increments the counter for the given key.
func (c *SafeCounter) Inc(key string) {
	c.mux.Lock()
	c.v[key]++
	c.mux.Unlock()
}

func (c *SafeCounter) Value(key string) int {
	c.mux.Lock()
	defer c.mux.Unlock()
	return c.v[key]
}

// end of samples
<commit_msg>[go][tour] Finished "Web Crawler"<commit_after>package main

import (
	"fmt"
	"github.com/hermantai/samples/go/commonutil"
	"golang.org/x/tour/tree"
	"sync"
	"time"
)

func main() {
	commonutil.PrintSection("Goroutines")
	goroutinesSample()

	commonutil.PrintSection("Channels")
	channelsSample()

	commonutil.PrintSection("Equivalent binary tree")
	fmt.Println(Same(tree.New(1), tree.New(1)))
	fmt.Println(Same(tree.New(1), tree.New(2)))

	commonutil.PrintSection("Locks")
	locksSample()

	commonutil.PrintSection("exercise: web crawler")
	webcrawlerSample()
}

// end of main

func goroutinesSample() {
	go say("world")
	say("hello")
}

func say(s string) {
	for i := 0; i < 5; i++ {
		time.Sleep(100 * time.Millisecond)
		fmt.Println(s)
	}
}

func channelsSample() {
	s := []int{7, 2, 8, -9, 4, 0}

	c := make(chan int)
	go sum(s[:len(s)/2], c)
	go sum(s[len(s)/2:], c)

	x, y := <-c, <-c

	fmt.Println(x, y, x+y)

	commonutil.PrintSubsection("Buffered channels")
	ch := make(chan int, 2)
	ch <- 3
	ch <- 4
	fmt.Println(<-ch, <-ch)

	ch2 := make(chan int, 10)
	go fibonacci(cap(ch2), ch2)
	for i := range ch2 {
		fmt.Println(i)
	}

	commonutil.PrintSubsection("Select")
	ch3 := make(chan int)
	quit := make(chan int)
	go func() {
		for i := 0; i < 10; i++ {
			fmt.Println(<-ch3)
		}
		quit <- 0
	}()
	fibWithSelect(ch3, quit)

	commonutil.PrintSubsection("default selection")
	tick := time.Tick(100 * time.Millisecond)
	boom := time.After(500 * time.Millisecond)
	for {
		select {
		case <-tick:
			fmt.Println("tick.")
		case <-boom:
			fmt.Println("BOOM!")
			return
		default:
			fmt.Println("    .")
			time.Sleep(50 * time.Millisecond)
		}
	}
}

func sum(s []int, c chan int) {
	sum := 0
	for _, v := range s {
		sum += v
	}
	c <- sum
}

func fibonacci(n int, c chan int) {
	x, y := 0, 1
	for i := 0; i < n; i++ {
		c <- x
		x, y = y, x+y
	}
	close(c)
}

func fibWithSelect(c, quit chan int) {
	x, y := 0, 1
	for {
		select {
		case c <- x:
			x, y = y, x+y
		case <-quit:
			fmt.Println("quit")
			return
		}
	}
}

// Walk walks the tree t sending all values
// from the tree to the channel ch.
func Walk(t *tree.Tree, ch chan int) {
	if t == nil {
		return
	}
	Walk(t.Left, ch)
	ch <- t.Value
	Walk(t.Right, ch)
}

// Same determines whether the trees
// t1 and t2 contain the same values.
func Same(t1, t2 *tree.Tree) bool {
	ch1, ch2 := make(chan int), make(chan int)
	go func() {
		Walk(t1, ch1)
		close(ch1)
	}()
	go func() {
		Walk(t2, ch2)
		close(ch2)
	}()

	for {
		x, ok := <-ch1
		if !ok {
			_, ok := <-ch2
			return !ok
		}

		if y, ok := <-ch2; !ok || x != y {
			return false
		}
	}
}

func locksSample() {
	c := SafeCounter{v: make(map[string]int)}
	for i := 0; i < 1000; i++ {
		go c.Inc("somekey")
	}

	time.Sleep(time.Second)
	fmt.Println(c.Value("somekey"))
}

// SafeCounter is safe to use concurrently.
type SafeCounter struct {
	v   map[string]int
	mux sync.Mutex
}

// Inc increments the counter for the given key.
func (c *SafeCounter) Inc(key string) {
	c.mux.Lock()
	c.v[key]++
	c.mux.Unlock()
}

func (c *SafeCounter) Value(key string) int {
	c.mux.Lock()
	defer c.mux.Unlock()
	return c.v[key]
}

func webcrawlerSample() {
	Crawl("https://golang.org/", 4, fetcher)
}

type Fetcher interface {
	// Fetch returns the body of URL and
	// a slice of URLs found on that page.
	Fetch(url string) (body string, urls []string, err error)
}

// Crawl uses fetcher to recursively crawl
// pages starting with url, to a maximum of depth.
func Crawl(url string, depth int, fetcher Fetcher) {
	fetchedUrlsLock.Lock()
	if _, fetched := attemptedUrls[url]; fetched {
		fetchedUrlsLock.Unlock()
		return
	}
	attemptedUrls[url] = struct{}{}
	fetchedUrlsLock.Unlock()

	if depth <= 0 {
		return
	}
	body, urls, err := fetcher.Fetch(url)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Printf("found: %s %q\n", url, body)
	var wg sync.WaitGroup
	for _, u := range urls {
		wg.Add(1)
		go func(url string) {
			defer wg.Done()
			Crawl(url, depth-1, fetcher)
		}(u)
	}
	wg.Wait()
	return
}

// fakeFetcher is Fetcher that returns canned results.
type fakeFetcher map[string]*fakeResult

type fakeResult struct {
	body string
	urls []string
}

func (f fakeFetcher) Fetch(url string) (string, []string, error) {
	if res, ok := f[url]; ok {
		return res.body, res.urls, nil
	}
	return "", nil, fmt.Errorf("not found: %s", url)
}

// fetcher is a populated fakeFetcher.
var fetcher = fakeFetcher{
	"https://golang.org/": &fakeResult{
		"The Go Programming Language",
		[]string{
			"https://golang.org/pkg/",
			"https://golang.org/cmd/",
		},
	},
	"https://golang.org/pkg/": &fakeResult{
		"Packages",
		[]string{
			"https://golang.org/",
			"https://golang.org/cmd/",
			"https://golang.org/pkg/fmt/",
			"https://golang.org/pkg/os/",
		},
	},
	"https://golang.org/pkg/fmt/": &fakeResult{
		"Package fmt",
		[]string{
			"https://golang.org/",
			"https://golang.org/pkg/",
		},
	},
	"https://golang.org/pkg/os/": &fakeResult{
		"Package os",
		[]string{
			"https://golang.org/",
			"https://golang.org/pkg/",
		},
	},
}

var attemptedUrls = make(map[string]interface{})
var fetchedUrlsLock sync.Mutex

// end of samples
<|endoftext|>