<commit_before>package main

import (
    "log"
    "github.com/fsouza/go-dockerclient"
    "io"
    "fmt"
    "bufio"
    "strings"
    "time"
    "bytes"
)

//DockerTaskDefinition is the varaible setting requests set by the user
type DockerTaskDefinition struct {
    ImageName   string  `json:"ImageName"`
    MacAddress  string  `json:"MacAddress"`
}

//AWSDOCKER is a dockerobj. It is identified by an image containerName
type AWSDOCKER struct {
    containerName        string
    taskArn              string
    timeout              time.Duration
    dockerClient         *docker.Client
    eventsCh             chan *docker.APIEvents
    containerArgs        string
    dockerTaskDefinition DockerTaskDefinition
}

func (dockerobj *AWSDOCKER) createDockerContainer(args []string, env []string, attachStdout bool) (string, error) {
    var taskPayloadEnv []string
    taskPayloadEnv = append(taskPayloadEnv, fmt.Sprintf("TASK_PAYLOAD=%s", dockerobj.containerArgs))

    dockerConfig := docker.Config{
        Env: taskPayloadEnv,
        Image: dockerobj.dockerTaskDefinition.ImageName,
        AttachStdout: attachStdout,
        AttachStderr: attachStdout,
        MacAddress: dockerobj.dockerTaskDefinition.MacAddress,
    }
    copts := docker.CreateContainerOptions{Name: dockerobj.containerName, Config: &dockerConfig}
    log.Printf("Create container for image container name: %s\n",dockerobj.dockerTaskDefinition.ImageName)
    container, err := dockerobj.dockerClient.CreateContainer(copts)
    if err != nil {
        return "", err
    }
    log.Printf("Created container container name: %s\n", container.ID)
    return container.ID, err
}

func (dockerobj *AWSDOCKER) deployImage(args []string, env []string, reader io.Reader) error {
    outputbuf := bytes.NewBuffer(nil)
    result := strings.Split(dockerobj.dockerTaskDefinition.ImageName, ":")
    opts := docker.PullImageOptions{
        Repository: result[0],
        Tag: result[1],
    }
    auth := docker.AuthConfiguration{}

    if err := dockerobj.dockerClient.PullImage(opts, auth); err != nil {
        log.Printf("Error building images: %s", err)
        log.Printf("Image Output:\n********************\n%s\n********************", outputbuf.String())
        return err
    }

    log.Printf("Created image: %s", dockerobj.dockerTaskDefinition.ImageName)

    return nil
}

//Deploy use the reader containing targz to create a docker image
//for docker inputbuf is tar reader ready for use by docker.Client
//the stream from end dockerClient to peer could directly be this tar stream
//talk to docker daemon using docker Client and build the image
func (dockerobj *AWSDOCKER) Deploy(args []string, env []string, reader io.Reader) error {
    if err := dockerobj.deployImage(args, env, reader); err != nil {
        return err
    }
    return nil
}

//BuildSpecFactory Should be removed
type BuildSpecFactory func() (io.Reader, error)

func (dockerobj *AWSDOCKER) stopInternal(id string, timeout uint, dontkill bool, dontremove bool) error {

    err := dockerobj.dockerClient.StopContainer(id, timeout)
    if err != nil {
        log.Printf("Stop container %s(%s)", id, err)
    } else {
        log.Printf("Stopped container %s", id)
    }
    if !dontkill {
        err = dockerobj.dockerClient.KillContainer(docker.KillContainerOptions{ID: id})
        if err != nil {
            log.Printf("Kill container %s (%s)", id, err)
        } else {
            log.Printf("Killed container %s", id)
        }
    }
    if !dontremove {
        err = dockerobj.dockerClient.RemoveContainer(docker.RemoveContainerOptions{ID: id, Force: true})
        if err != nil {
            log.Printf("Remove container %s (%s)", id, err)
        } else {
            log.Printf("Removed container %s", id)
        }
    }
    return err
}

//Start starts a container using a previously created docker image
func (dockerobj *AWSDOCKER) Start(messageBody *string, args []string, env []string, builder BuildSpecFactory,  messageID *string) error {

    attachStdout := true

    //start ECS Task

    //stop,force remove if necessary
    log.Printf("Cleanup image containerName %s", dockerobj.containerName)

    dockerobj.stopInternal(dockerobj.containerName, 0, false, false)

    log.Printf("Start container %s", dockerobj.containerName)
    containerID, err := dockerobj.createDockerContainer(args, env, attachStdout)
    dockerobj.taskArn = containerID
    if err != nil {
        //if image not found try to create image and retry
        if err == docker.ErrNoSuchImage {
            if builder == nil {
                log.Printf("start-could not find image ...attempt to recreate image %s", err)
                var err1 error
                //reader, err1 := builder()
                //if err1 != nil {
                //    log.Printf("Error creating image builder: %s", err1)
                //}

                if err1 = dockerobj.deployImage(args, env, nil); err1 != nil {
                    return err1
                }

                log.Printf("start-recreated image successfully")
                if containerID, err1 = dockerobj.createDockerContainer(args, env, attachStdout); err1 != nil {
                    log.Printf("start-could not recreate container post recreate image: %s", err1)
                    return err1
                }
            } else {
                log.Printf("start-could not find image: %s", err)
                return err
            }
        } else {
            log.Printf("start-could not recreate container: %s", err)
            return err
        }
    }

    if attachStdout {
        // Launch a few go-threads to manage output streams from the container.
        // They will be automatically destroyed when the container exits
        attached := make(chan struct{})
        r, w := io.Pipe()

        go func() {
            // AttachToContainer will fire off a message on the "attached" channel once the
            // attachment completes, and then block until the container is terminated.
            // The returned error is not used outside the scope of this function. Assign the
            // error to a local variable to prevent clobbering the function variable 'err'.
            err := dockerobj.dockerClient.AttachToContainer(docker.AttachToContainerOptions{
                Container:    containerID,
                OutputStream: w,
                ErrorStream:  w,
                Logs:         true,
                Stdout:       true,
                Stderr:       true,
                Stream:       true,
                Success:      attached,
            })

            // If we get here, the container has terminated.  Send a signal on the pipe
            // so that downstream may clean up appropriately
            _ = w.CloseWithError(err)
        }()

        go func() {
            // Block here until the attachment completes or we timeout
            select {
            case <-attached:
                // successful attach
            case <-time.After(10 * time.Second):
                log.Printf("Timeout while attaching to IO channel in container %s", containerID)
                return
            }

            // Acknowledge the attachment?  This was included in the gist I followed
            // (http://bit.ly/2jBrCtM).  Not sure it's actually needed but it doesn't
            // appear to hurt anything.
            attached <- struct{}{}

            // Establish a buffer for our IO channel so that we may do readline-style
            // ingestion of the IO, one log entry per line
            is := bufio.NewReader(r)

            for {
                // Loop forever dumping lines of text into the containerLogger
                // until the pipe is closed
                line, err2 := is.ReadString('\n')
                if err2 != nil {
                    switch err2 {
                    case io.EOF:
                        log.Printf("Container %s has closed its IO channel", containerID)
                    default:
                        log.Printf("Error reading container output: %s", err2)
                    }

                    return
                }

                log.Printf(line)
            }
        }()
    }

    err = dockerobj.dockerClient.StartContainer(containerID, nil)
    if err != nil {
        log.Printf("start-could not start container: %s", err)
        return err
    }

    log.Printf("Started container %s", containerID)
    return nil
}

//Stop stops a running chaincode
func (dockerobj *AWSDOCKER) Stop(id string, timeout uint, dontkill bool, dontremove bool) error {

    id = strings.Replace(id, ":", "_", -1)
    err := dockerobj.stopInternal(id, timeout, dontkill, dontremove)

    return err
}

//Destroy destroys an image
func (dockerobj *AWSDOCKER) Destroy(id string, force bool, noprune bool) error {
    id = strings.Replace(id, ":", "_", -1)

    err := dockerobj.dockerClient.RemoveImageExtended(id, docker.RemoveImageOptions{Force: force, NoPrune: noprune})

    if err != nil {
        log.Printf("error while destroying image: %s", err)
    } else {
        log.Printf("Destroyed image %s", id)
    }

    return err
}

func (dockerobj AWSDOCKER) execute(handler MessageHandler) {
    handler.initialize()
    if handler.receive() {
        dockerobj.dockerobjTimeoutHelper(handler)
    }
}

func (dockerobj *AWSDOCKER) dockerobjTimeoutHelper(handler MessageHandler) {
    ch := make(chan error)
    go func() {
        ch <- dockerobj.executionHelper(handler.body(), handler.id())
    }()
    select {
    case err := <-ch:
        if err != nil {
            log.Printf("E: %s %s", dockerobj.containerName, err.Error())
            handler.failure(err)
        } else {
            log.Printf("I: %s finished successfully", dockerobj.containerName)
            handler.success()
        }
    case <-time.After(dockerobj.timeout):
        err := fmt.Errorf("%s timed out after %f seconds", dockerobj.containerName, dockerobj.timeout.Seconds())
        log.Println(err)
        handler.failure(err)
    }
}

func (dockerobj *AWSDOCKER) executionHelper(messageBody *string, messageID *string) error {
    var err error

    args := make([]string, 1)
    env := make([]string, 1)
    env = append(env, *messageBody)

    //taskArn, err = dockerobj.startECSTask(messageBody, messageID)
    //dockerobj.taskArn = taskArn

    err = dockerobj.Start(messageBody, args, env, nil, messageID)
    if err != nil {
        return err
    }
    err = dockerobj.monitorDocker()
    if err != nil {
        return err
    }
    return nil
}

func (dockerobj *AWSDOCKER) monitorDocker() error {
    dockerobj.addListener()
    // Monitor docker events for sibling Projector task
    status, err := dockerobj.listenForDie()
    if err != nil {
        return err
    }

    if status == "0" {
        // status is die
        log.Printf("[INFO] Execution completed successfully")
        dockerobj.success()
        return nil
    }
    // non-zero exit
    log.Printf("[ERROR] Execution completed with non-zero exit status")
    err = fmt.Errorf("%s died with non-zero exit status (exit code %s)", dockerobj.containerName, status)
    dockerobj.failure()
    return err

}

func (dockerobj *AWSDOCKER) listenForDie() (exitCode string, err error) {
    log.Printf("[INFO] Monitoring Docker events.")
    log.Printf("[DEBUG] %+v\n", dockerobj)
    duration := getTimeout()
    timeout := time.After(duration)
    defer dockerobj.removeListener()
    for {
        select {
        case msg := <-dockerobj.eventsCh:
            if msg != nil {
                matched := msg.Actor.ID == dockerobj.taskArn
                if matched {
                    log.Printf("[DEBUG] %+v\n", msg)
                    switch msg.Action {
                    case "die":
                        log.Printf("[INFO] Container die event")
                        return msg.Actor.Attributes["exitCode"], nil
                    }
                }
            }
        case <-timeout:
            log.Printf("[INFO] Instance timeout reached.")
            err := fmt.Errorf("Docker container %s timed out after %f seconds", dockerobj.containerName, duration.Seconds())
            return "timeout", err
        }
    }
}

func (dockerobj *AWSDOCKER) connect(dockerEndpointPath string) {
    log.Printf("[INFO] Connecting to Docker API.")
    endpoint := dockerEndpointPath
    client, err := docker.NewClient(endpoint)
    if err != nil {
        panic(err)
    }
    dockerobj.dockerClient = client
    dockerobj.eventsCh = make(chan *docker.APIEvents)
}

func (dockerobj *AWSDOCKER) addListener() {
    err := dockerobj.dockerClient.AddEventListener(dockerobj.eventsCh)
    if err != nil {
        log.Fatal(err)
    }
}

func (dockerobj *AWSDOCKER) removeListener() {
    err := dockerobj.dockerClient.RemoveEventListener(dockerobj.eventsCh)
    if err != nil {
        log.Fatal(err)
    }
}

func (dockerobj *AWSDOCKER) success() {}
func (dockerobj *AWSDOCKER) failure() {}<commit_msg>Reformatted files, and used messageBody in createDockerContainer<commit_after>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"log"
	"strings"
	"time"

	"github.com/fsouza/go-dockerclient"
)

//DockerTaskDefinition is the varaible setting requests set by the user
type DockerTaskDefinition struct {
	ImageName  string `json:"ImageName"`
	MacAddress string `json:"MacAddress"`
}

//AWSDOCKER is a dockerobj. It is identified by an image containerName
type AWSDOCKER struct {
	containerName        string
	taskArn              string
	timeout              time.Duration
	dockerClient         *docker.Client
	eventsCh             chan *docker.APIEvents
	containerArgs        string
	dockerTaskDefinition DockerTaskDefinition
}

func (dockerobj *AWSDOCKER) createDockerContainer(messageBody *string, args []string, env []string, attachStdout bool) (string, error) {
	var taskPayloadEnv []string
	taskPayloadEnv = append(taskPayloadEnv, fmt.Sprintf("TASK_PAYLOAD=%s", *messageBody))

	dockerConfig := docker.Config{
		Env:          taskPayloadEnv,
		Image:        dockerobj.dockerTaskDefinition.ImageName,
		AttachStdout: attachStdout,
		AttachStderr: attachStdout,
		MacAddress:   dockerobj.dockerTaskDefinition.MacAddress,
	}
	copts := docker.CreateContainerOptions{Name: dockerobj.containerName, Config: &dockerConfig}
	log.Printf("Create container for image container name: %s\n", dockerobj.dockerTaskDefinition.ImageName)
	container, err := dockerobj.dockerClient.CreateContainer(copts)
	if err != nil {
		return "", err
	}
	log.Printf("Created container container name: %s\n", container.ID)
	return container.ID, err
}

func (dockerobj *AWSDOCKER) deployImage(args []string, env []string, reader io.Reader) error {
	outputbuf := bytes.NewBuffer(nil)
	result := strings.Split(dockerobj.dockerTaskDefinition.ImageName, ":")
	opts := docker.PullImageOptions{
		Repository: result[0],
		Tag:        result[1],
	}
	auth := docker.AuthConfiguration{}

	if err := dockerobj.dockerClient.PullImage(opts, auth); err != nil {
		log.Printf("Error building images: %s", err)
		log.Printf("Image Output:\n********************\n%s\n********************", outputbuf.String())
		return err
	}

	log.Printf("Created image: %s", dockerobj.dockerTaskDefinition.ImageName)

	return nil
}

//Deploy use the reader containing targz to create a docker image
//for docker inputbuf is tar reader ready for use by docker.Client
//the stream from end dockerClient to peer could directly be this tar stream
//talk to docker daemon using docker Client and build the image
func (dockerobj *AWSDOCKER) Deploy(args []string, env []string, reader io.Reader) error {
	if err := dockerobj.deployImage(args, env, reader); err != nil {
		return err
	}
	return nil
}

//BuildSpecFactory Should be removed
type BuildSpecFactory func() (io.Reader, error)

func (dockerobj *AWSDOCKER) stopInternal(id string, timeout uint, dontkill bool, dontremove bool) error {

	err := dockerobj.dockerClient.StopContainer(id, timeout)
	if err != nil {
		log.Printf("Stop container %s(%s)", id, err)
	} else {
		log.Printf("Stopped container %s", id)
	}
	if !dontkill {
		err = dockerobj.dockerClient.KillContainer(docker.KillContainerOptions{ID: id})
		if err != nil {
			log.Printf("Kill container %s (%s)", id, err)
		} else {
			log.Printf("Killed container %s", id)
		}
	}
	if !dontremove {
		err = dockerobj.dockerClient.RemoveContainer(docker.RemoveContainerOptions{ID: id, Force: true})
		if err != nil {
			log.Printf("Remove container %s (%s)", id, err)
		} else {
			log.Printf("Removed container %s", id)
		}
	}
	return err
}

//Start starts a container using a previously created docker image
func (dockerobj *AWSDOCKER) Start(messageBody *string, args []string, env []string, builder BuildSpecFactory, messageID *string) error {

	attachStdout := true

	//start ECS Task

	//stop,force remove if necessary
	log.Printf("Cleanup image containerName %s", dockerobj.containerName)

	dockerobj.stopInternal(dockerobj.containerName, 0, false, false)

	log.Printf("Start container %s", dockerobj.containerName)
	containerID, err := dockerobj.createDockerContainer(messageBody, args, env, attachStdout)
	dockerobj.taskArn = containerID
	if err != nil {
		//if image not found try to create image and retry
		if err == docker.ErrNoSuchImage {
			if builder == nil {
				log.Printf("start-could not find image ...attempt to recreate image %s", err)
				var err1 error
				//reader, err1 := builder()
				//if err1 != nil {
				//    log.Printf("Error creating image builder: %s", err1)
				//}

				if err1 = dockerobj.deployImage(args, env, nil); err1 != nil {
					return err1
				}

				log.Printf("start-recreated image successfully")
				if containerID, err1 = dockerobj.createDockerContainer(messageBody, args, env, attachStdout); err1 != nil {
					log.Printf("start-could not recreate container post recreate image: %s", err1)
					return err1
				}
			} else {
				log.Printf("start-could not find image: %s", err)
				return err
			}
		} else {
			log.Printf("start-could not recreate container: %s", err)
			return err
		}
	}

	if attachStdout {
		// Launch a few go-threads to manage output streams from the container.
		// They will be automatically destroyed when the container exits
		attached := make(chan struct{})
		r, w := io.Pipe()

		go func() {
			// AttachToContainer will fire off a message on the "attached" channel once the
			// attachment completes, and then block until the container is terminated.
			// The returned error is not used outside the scope of this function. Assign the
			// error to a local variable to prevent clobbering the function variable 'err'.
			err := dockerobj.dockerClient.AttachToContainer(docker.AttachToContainerOptions{
				Container:    containerID,
				OutputStream: w,
				ErrorStream:  w,
				Logs:         true,
				Stdout:       true,
				Stderr:       true,
				Stream:       true,
				Success:      attached,
			})

			// If we get here, the container has terminated.  Send a signal on the pipe
			// so that downstream may clean up appropriately
			_ = w.CloseWithError(err)
		}()

		go func() {
			// Block here until the attachment completes or we timeout
			select {
			case <-attached:
				// successful attach
			case <-time.After(10 * time.Second):
				log.Printf("Timeout while attaching to IO channel in container %s", containerID)
				return
			}

			// Acknowledge the attachment?  This was included in the gist I followed
			// (http://bit.ly/2jBrCtM).  Not sure it's actually needed but it doesn't
			// appear to hurt anything.
			attached <- struct{}{}

			// Establish a buffer for our IO channel so that we may do readline-style
			// ingestion of the IO, one log entry per line
			is := bufio.NewReader(r)

			for {
				// Loop forever dumping lines of text into the containerLogger
				// until the pipe is closed
				line, err2 := is.ReadString('\n')
				if err2 != nil {
					switch err2 {
					case io.EOF:
						log.Printf("Container %s has closed its IO channel", containerID)
					default:
						log.Printf("Error reading container output: %s", err2)
					}

					return
				}

				log.Printf(line)
			}
		}()
	}

	err = dockerobj.dockerClient.StartContainer(containerID, nil)
	if err != nil {
		log.Printf("start-could not start container: %s", err)
		return err
	}

	log.Printf("Started container %s", containerID)
	return nil
}

//Stop stops a running chaincode
func (dockerobj *AWSDOCKER) Stop(id string, timeout uint, dontkill bool, dontremove bool) error {

	id = strings.Replace(id, ":", "_", -1)
	err := dockerobj.stopInternal(id, timeout, dontkill, dontremove)

	return err
}

//Destroy destroys an image
func (dockerobj *AWSDOCKER) Destroy(id string, force bool, noprune bool) error {
	id = strings.Replace(id, ":", "_", -1)

	err := dockerobj.dockerClient.RemoveImageExtended(id, docker.RemoveImageOptions{Force: force, NoPrune: noprune})

	if err != nil {
		log.Printf("error while destroying image: %s", err)
	} else {
		log.Printf("Destroyed image %s", id)
	}

	return err
}

func (dockerobj AWSDOCKER) execute(handler MessageHandler) {
	handler.initialize()
	if handler.receive() {
		dockerobj.dockerobjTimeoutHelper(handler)
	}
}

func (dockerobj *AWSDOCKER) dockerobjTimeoutHelper(handler MessageHandler) {
	ch := make(chan error)
	go func() {
		ch <- dockerobj.executionHelper(handler.body(), handler.id())
	}()
	select {
	case err := <-ch:
		if err != nil {
			log.Printf("E: %s %s", dockerobj.containerName, err.Error())
			handler.failure(err)
		} else {
			log.Printf("I: %s finished successfully", dockerobj.containerName)
			handler.success()
		}
	case <-time.After(dockerobj.timeout):
		err := fmt.Errorf("%s timed out after %f seconds", dockerobj.containerName, dockerobj.timeout.Seconds())
		log.Println(err)
		handler.failure(err)
	}
}

func (dockerobj *AWSDOCKER) executionHelper(messageBody *string, messageID *string) error {
	var err error

	args := make([]string, 1)
	env := make([]string, 1)
	env = append(env, *messageBody)
	fmt.Println(env)

	//taskArn, err = dockerobj.startECSTask(messageBody, messageID)
	//dockerobj.taskArn = taskArn

	err = dockerobj.Start(messageBody, args, env, nil, messageID)
	if err != nil {
		return err
	}
	err = dockerobj.monitorDocker()
	if err != nil {
		return err
	}
	return nil
}

func (dockerobj *AWSDOCKER) monitorDocker() error {
	dockerobj.addListener()
	// Monitor docker events for sibling Projector task
	status, err := dockerobj.listenForDie()
	if err != nil {
		return err
	}

	if status == "0" {
		// status is die
		log.Printf("[INFO] Execution completed successfully")
		dockerobj.success()
		return nil
	}
	// non-zero exit
	log.Printf("[ERROR] Execution completed with non-zero exit status")
	err = fmt.Errorf("%s died with non-zero exit status (exit code %s)", dockerobj.containerName, status)
	dockerobj.failure()
	return err

}

func (dockerobj *AWSDOCKER) listenForDie() (exitCode string, err error) {
	log.Printf("[INFO] Monitoring Docker events.")
	log.Printf("[DEBUG] %+v\n", dockerobj)
	duration := getTimeout()
	timeout := time.After(duration)
	defer dockerobj.removeListener()
	for {
		select {
		case msg := <-dockerobj.eventsCh:
			if msg != nil {
				matched := msg.Actor.ID == dockerobj.taskArn
				if matched {
					log.Printf("[DEBUG] %+v\n", msg)
					switch msg.Action {
					case "die":
						log.Printf("[INFO] Container die event")
						return msg.Actor.Attributes["exitCode"], nil
					}
				}
			}
		case <-timeout:
			log.Printf("[INFO] Instance timeout reached.")
			err := fmt.Errorf("Docker container %s timed out after %f seconds", dockerobj.containerName, duration.Seconds())
			return "timeout", err
		}
	}
}

func (dockerobj *AWSDOCKER) connect(dockerEndpointPath string) {
	log.Printf("[INFO] Connecting to Docker API.")
	endpoint := dockerEndpointPath
	client, err := docker.NewClient(endpoint)
	if err != nil {
		panic(err)
	}
	dockerobj.dockerClient = client
	dockerobj.eventsCh = make(chan *docker.APIEvents)
}

func (dockerobj *AWSDOCKER) addListener() {
	err := dockerobj.dockerClient.AddEventListener(dockerobj.eventsCh)
	if err != nil {
		log.Fatal(err)
	}
}

func (dockerobj *AWSDOCKER) removeListener() {
	err := dockerobj.dockerClient.RemoveEventListener(dockerobj.eventsCh)
	if err != nil {
		log.Fatal(err)
	}
}

func (dockerobj *AWSDOCKER) success() {}
func (dockerobj *AWSDOCKER) failure() {}
<|endoftext|>


<commit_before>// Copyright 2022 Google LLC
//
// 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 mixed_oc_cli_origin_support_test implements GNMI 1.12 from go/wbb:vendor-testplan
package mixed_oc_cli_origin_support_test

import (
	"context"
	"fmt"
	"testing"

	"github.com/openconfig/featureprofiles/internal/fptest"
	gpb "github.com/openconfig/gnmi/proto/gnmi"
	"github.com/openconfig/ondatra"
	"github.com/openconfig/ygot/ygot"
)

func TestMain(m *testing.M) {
	fptest.RunTests(m)
}

func interfaceDescriptionCLI(dp *ondatra.Port, desc string) string {
	switch dp.Device().Vendor() {
	case ondatra.ARISTA, ondatra.CISCO:
		const tmpl = `
interface %s
  description %s
`
		return fmt.Sprintf(tmpl, dp.Name(), desc)
	}
	return ""
}

func buildOCUpdate(path *gpb.Path, value string) *gpb.Update {
	if len(path.GetElem()) == 0 || path.GetElem()[0].GetName() != "meta" {
		path.Origin = "openconfig"
	}
	jsonVal, _ := ygot.Marshal7951(ygot.String(value))
	update := &gpb.Update{
		Path: path,
		Val:  &gpb.TypedValue{Value: &gpb.TypedValue_JsonIetfVal{JsonIetfVal: jsonVal}},
	}
	return update
}

func buildCLIUpdate(value string) *gpb.Update {
	update := &gpb.Update{
		Path: &gpb.Path{
			Origin: "cli",
			Elem:   []*gpb.PathElem{},
		},
		Val: &gpb.TypedValue{
			Value: &gpb.TypedValue_AsciiVal{
				AsciiVal: value,
			},
		},
	}
	return update
}

// TestCLIBeforeOpenConfig pushes overlapping mixed SetRequest specifying CLI before OpenConfig for DUT port-1.
func TestCLIBeforeOpenConfig(t *testing.T) {
	dut := ondatra.DUT(t, "dut")
	dp := dut.Port(t, "port1")

	// `origin: "cli"` - containing vendor configuration.
	intfConfig := interfaceDescriptionCLI(dp, "from cli")
	if intfConfig == "" {
		t.Fatalf("Please add vendor support for %v", dut.Vendor())
	}
	t.Logf("Building the CLI config:\n%s", intfConfig)

	// `origin: ""` (openconfig, default origin) setting the DUT port-1
	//  string value at `/interfaces/interface/config/description` to `"from oc"`.
	resolvedPath := dut.Config().Interface(dp.Name()).Description()
	path, _, errs := ygot.ResolvePath(resolvedPath)
	if errs != nil {
		t.Fatalf("Could not resolve path: %v", errs)
	}

	gpbSetRequest := &gpb.SetRequest{
		Update: []*gpb.Update{
			buildCLIUpdate(intfConfig),
			buildOCUpdate(path, "from oc"),
		},
	}
	t.Log("gnmiClient Set both CLI and OpenConfig modelled config")
	t.Log(gpbSetRequest)

	gnmiClient := dut.RawAPIs().GNMI().Default(t)
	response, err := gnmiClient.Set(context.Background(), gpbSetRequest)
	if err != nil {
		t.Fatalf("gnmiClient.Set() with unexpected error: %v", err)
	}
	t.Log("gnmiClient Set Response for CLI and OpenConfig modelled config")
	t.Log(response)

	// Validate that DUT port-1 description is `"from oc"`
	got := dut.Telemetry().Interface(dp.Name()).Description().Get(t)
	want := "from oc"
	if got != want {
		t.Errorf("Get(DUT port description): got %v, want %v", got, want)
	}
}

// TestOpenConfigBeforeCLI pushes overlapping mixed SetRequest specifying OpenConfig before CLI for DUT port-1.
func TestOpenConfigBeforeCLI(t *testing.T) {
	dut := ondatra.DUT(t, "dut")
	dp := dut.Port(t, "port1")

	// `origin: ""` (openconfig, default origin) setting the DUT port-1
	//  string value at `/interfaces/interface/config/description` to `"from oc"`.
	resolvedPath := dut.Config().Interface(dp.Name()).Description()
	path, _, errs := ygot.ResolvePath(resolvedPath)
	if errs != nil {
		t.Fatalf("Could not resolve path: %v", errs)
	}

	// `origin: "cli"` - containing vendor configuration.
	intfConfig := interfaceDescriptionCLI(dp, "from cli")
	if intfConfig == "" {
		t.Fatalf("Please add vendor support for %v", dut.Vendor())
	}
	t.Logf("Building the CLI config:\n%s", intfConfig)

	gpbSetRequest := &gpb.SetRequest{
		Update: []*gpb.Update{
			buildOCUpdate(path, "from oc"),
			buildCLIUpdate(intfConfig),
		},
	}
	t.Log("gnmiClient Set both CLI and OpenConfig modelled config")
	t.Log(gpbSetRequest)

	gnmiClient := dut.RawAPIs().GNMI().Default(t)
	response, err := gnmiClient.Set(context.Background(), gpbSetRequest)
	if err != nil {
		t.Fatalf("gnmiClient.Set() with unexpected error: %v", err)
	}
	t.Log("gnmiClient Set Response for CLI and OpenConfig modelled config")
	t.Log(response)

	// Validate that DUT port-1 description is `"from cli"`
	got := dut.Telemetry().Interface(dp.Name()).Description().Get(t)
	want := "from cli"
	if got != want {
		t.Errorf("Get(DUT port description): got %v, want %v", got, want)
	}
}

// Push non-overlapping mixed SetRequest specifying CLI for DUT port-1 and
// OpenConfig for DUT port-2
func TestMixedOriginOCCLIConfig(t *testing.T) {
	dut := ondatra.DUT(t, "dut")
	dp1 := dut.Port(t, "port1")
	dp2 := dut.Port(t, "port2")
	gnmiClient := dut.RawAPIs().GNMI().Default(t)
	// `origin: "cli"` - containing vendor configuration.
	intf1Config := interfaceDescriptionCLI(dp1, "foo1")
	if intf1Config == "" {
		t.Fatalf("Please add vendor support for %v", dut.Vendor())
	}
	t.Logf("Building the CLI config:\n%s", intf1Config)

	// `origin: ""` (openconfig, default origin) setting the DUT port-2
	//  string value at `/interfaces/interface/config/description` to `"foo2"`.
	resolvedPath := dut.Config().Interface(dp2.Name()).Description()
	path, _, errs := ygot.ResolvePath(resolvedPath)

	gpbSetRequest := &gpb.SetRequest{
		Update: []*gpb.Update{
			buildCLIUpdate(intf1Config),
			buildOCUpdate(path, "foo2"),
		},
	}
	if errs != nil {
		t.Fatalf("Could not resolve path: %v", errs)
	}
	t.Log("gnmiClient Set both CLI and OpenConfig modelled config")
	t.Log(gpbSetRequest)

	response, err := gnmiClient.Set(context.Background(), gpbSetRequest)
	if err != nil {
		t.Fatalf("gnmiClient.Set() with unexpected error: %v", err)
	}
	t.Log("gnmiClient Set Response for CLI and OpenConfig modelled config")
	t.Log(response)

	// Validate that DUT port-1 and DUT port-2 description through telemetry.
	if got := dut.Telemetry().Interface(dp1.Name()).Description().Get(t); got != "foo1" {
		t.Errorf("Get(DUT port description): got %v, want %v", got, "foo1")
	}
	if got := dut.Telemetry().Interface(dp2.Name()).Description().Get(t); got != "foo2" {
		t.Errorf("Get(DUT port description): got %v, want %v", got, "foo2")
	}

}
<commit_msg>gNMI-1.12: add juniper cli command (#706)<commit_after>// Copyright 2022 Google LLC
//
// 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 mixed_oc_cli_origin_support_test implements GNMI 1.12 from go/wbb:vendor-testplan
package mixed_oc_cli_origin_support_test

import (
	"context"
	"fmt"
	"testing"

	"github.com/openconfig/featureprofiles/internal/fptest"
	gpb "github.com/openconfig/gnmi/proto/gnmi"
	"github.com/openconfig/ondatra"
	"github.com/openconfig/ygot/ygot"
)

func TestMain(m *testing.M) {
	fptest.RunTests(m)
}

func interfaceDescriptionCLI(dp *ondatra.Port, desc string) string {
	switch dp.Device().Vendor() {
	case ondatra.ARISTA, ondatra.CISCO:
		const tmpl = `
interface %s
  description %s
`
		return fmt.Sprintf(tmpl, dp.Name(), desc)
	case ondatra.JUNIPER:
		const tmpl = `
interfaces {
    %s {
        description "%s"
    }
}
`
		return fmt.Sprintf(tmpl, dp.Name(), desc)
	}
	return ""
}

func buildOCUpdate(path *gpb.Path, value string) *gpb.Update {
	if len(path.GetElem()) == 0 || path.GetElem()[0].GetName() != "meta" {
		path.Origin = "openconfig"
	}
	jsonVal, _ := ygot.Marshal7951(ygot.String(value))
	update := &gpb.Update{
		Path: path,
		Val:  &gpb.TypedValue{Value: &gpb.TypedValue_JsonIetfVal{JsonIetfVal: jsonVal}},
	}
	return update
}

func buildCLIUpdate(value string) *gpb.Update {
	update := &gpb.Update{
		Path: &gpb.Path{
			Origin: "cli",
			Elem:   []*gpb.PathElem{},
		},
		Val: &gpb.TypedValue{
			Value: &gpb.TypedValue_AsciiVal{
				AsciiVal: value,
			},
		},
	}
	return update
}

// TestCLIBeforeOpenConfig pushes overlapping mixed SetRequest specifying CLI before OpenConfig for DUT port-1.
func TestCLIBeforeOpenConfig(t *testing.T) {
	dut := ondatra.DUT(t, "dut")
	dp := dut.Port(t, "port1")

	// `origin: "cli"` - containing vendor configuration.
	intfConfig := interfaceDescriptionCLI(dp, "from cli")
	if intfConfig == "" {
		t.Fatalf("Please add vendor support for %v", dut.Vendor())
	}
	t.Logf("Building the CLI config:\n%s", intfConfig)

	// `origin: ""` (openconfig, default origin) setting the DUT port-1
	//  string value at `/interfaces/interface/config/description` to `"from oc"`.
	resolvedPath := dut.Config().Interface(dp.Name()).Description()
	path, _, errs := ygot.ResolvePath(resolvedPath)
	if errs != nil {
		t.Fatalf("Could not resolve path: %v", errs)
	}

	gpbSetRequest := &gpb.SetRequest{
		Update: []*gpb.Update{
			buildCLIUpdate(intfConfig),
			buildOCUpdate(path, "from oc"),
		},
	}
	t.Log("gnmiClient Set both CLI and OpenConfig modelled config")
	t.Log(gpbSetRequest)

	gnmiClient := dut.RawAPIs().GNMI().Default(t)
	response, err := gnmiClient.Set(context.Background(), gpbSetRequest)
	if err != nil {
		t.Fatalf("gnmiClient.Set() with unexpected error: %v", err)
	}
	t.Log("gnmiClient Set Response for CLI and OpenConfig modelled config")
	t.Log(response)

	// Validate that DUT port-1 description is `"from oc"`
	got := dut.Telemetry().Interface(dp.Name()).Description().Get(t)
	want := "from oc"
	if got != want {
		t.Errorf("Get(DUT port description): got %v, want %v", got, want)
	}
}

// TestOpenConfigBeforeCLI pushes overlapping mixed SetRequest specifying OpenConfig before CLI for DUT port-1.
func TestOpenConfigBeforeCLI(t *testing.T) {
	dut := ondatra.DUT(t, "dut")
	dp := dut.Port(t, "port1")

	// `origin: ""` (openconfig, default origin) setting the DUT port-1
	//  string value at `/interfaces/interface/config/description` to `"from oc"`.
	resolvedPath := dut.Config().Interface(dp.Name()).Description()
	path, _, errs := ygot.ResolvePath(resolvedPath)
	if errs != nil {
		t.Fatalf("Could not resolve path: %v", errs)
	}

	// `origin: "cli"` - containing vendor configuration.
	intfConfig := interfaceDescriptionCLI(dp, "from cli")
	if intfConfig == "" {
		t.Fatalf("Please add vendor support for %v", dut.Vendor())
	}
	t.Logf("Building the CLI config:\n%s", intfConfig)

	gpbSetRequest := &gpb.SetRequest{
		Update: []*gpb.Update{
			buildOCUpdate(path, "from oc"),
			buildCLIUpdate(intfConfig),
		},
	}
	t.Log("gnmiClient Set both CLI and OpenConfig modelled config")
	t.Log(gpbSetRequest)

	gnmiClient := dut.RawAPIs().GNMI().Default(t)
	response, err := gnmiClient.Set(context.Background(), gpbSetRequest)
	if err != nil {
		t.Fatalf("gnmiClient.Set() with unexpected error: %v", err)
	}
	t.Log("gnmiClient Set Response for CLI and OpenConfig modelled config")
	t.Log(response)

	// Validate that DUT port-1 description is `"from cli"`
	got := dut.Telemetry().Interface(dp.Name()).Description().Get(t)
	want := "from cli"
	if got != want {
		t.Errorf("Get(DUT port description): got %v, want %v", got, want)
	}
}

// Push non-overlapping mixed SetRequest specifying CLI for DUT port-1 and
// OpenConfig for DUT port-2
func TestMixedOriginOCCLIConfig(t *testing.T) {
	dut := ondatra.DUT(t, "dut")
	dp1 := dut.Port(t, "port1")
	dp2 := dut.Port(t, "port2")
	gnmiClient := dut.RawAPIs().GNMI().Default(t)
	// `origin: "cli"` - containing vendor configuration.
	intf1Config := interfaceDescriptionCLI(dp1, "foo1")
	if intf1Config == "" {
		t.Fatalf("Please add vendor support for %v", dut.Vendor())
	}
	t.Logf("Building the CLI config:\n%s", intf1Config)

	// `origin: ""` (openconfig, default origin) setting the DUT port-2
	//  string value at `/interfaces/interface/config/description` to `"foo2"`.
	resolvedPath := dut.Config().Interface(dp2.Name()).Description()
	path, _, errs := ygot.ResolvePath(resolvedPath)

	gpbSetRequest := &gpb.SetRequest{
		Update: []*gpb.Update{
			buildCLIUpdate(intf1Config),
			buildOCUpdate(path, "foo2"),
		},
	}
	if errs != nil {
		t.Fatalf("Could not resolve path: %v", errs)
	}
	t.Log("gnmiClient Set both CLI and OpenConfig modelled config")
	t.Log(gpbSetRequest)

	response, err := gnmiClient.Set(context.Background(), gpbSetRequest)
	if err != nil {
		t.Fatalf("gnmiClient.Set() with unexpected error: %v", err)
	}
	t.Log("gnmiClient Set Response for CLI and OpenConfig modelled config")
	t.Log(response)

	// Validate that DUT port-1 and DUT port-2 description through telemetry.
	if got := dut.Telemetry().Interface(dp1.Name()).Description().Get(t); got != "foo1" {
		t.Errorf("Get(DUT port description): got %v, want %v", got, "foo1")
	}
	if got := dut.Telemetry().Interface(dp2.Name()).Description().Get(t); got != "foo2" {
		t.Errorf("Get(DUT port description): got %v, want %v", got, "foo2")
	}

}
<|endoftext|>


<commit_before>/*
Copyright (c) 2014 Ashley Jeffs

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
*/

package tunny

import (
	"testing"
	"time"
	"runtime"
)

func TestTimeout(t *testing.T) {
	outChan  := make(chan int, 3)

	pool, errPool := CreatePool(1, func(object interface{}) interface{} {
		time.Sleep(500 * time.Millisecond)
		return nil
	}).Open()

	if errPool != nil {
		t.Errorf("Error starting pool: ", errPool)
		return
	}

	defer pool.Close()

	before := time.Now()

	go func() {
		if _, err := pool.SendWorkTimed(200, nil); err == nil {
			t.Errorf("No timeout triggered thread one")
		} else {
			taken := ( time.Since(before) / time.Millisecond )
			if taken > 210 {
				t.Errorf("Time taken at thread one: ", taken, ", with error: ", err)
			}
		}
		outChan <- 1

		go func() {
			if _, err := pool.SendWork(nil); err == nil {
			} else {
				t.Errorf("Error at thread three: ", err)
			}
			outChan <- 1
		}()
	}()

	go func() {
		if _, err := pool.SendWorkTimed(200, nil); err == nil {
			t.Errorf("No timeout triggered thread two")
		} else {
			taken := ( time.Since(before) / time.Millisecond )
			if taken > 210 {
				t.Errorf("Time taken at thread two: ", taken, ", with error: ", err)
			}
		}
		outChan <- 1
	}()

	for i := 0; i < 3; i++ {
		<-outChan
	}
}

func TestTimeoutRequests(t *testing.T) {
	n_polls := 200
	outChan  := make(chan int, n_polls)

	pool, errPool := CreatePool(1, func(object interface{}) interface{} {
		time.Sleep(time.Millisecond)
		return nil
	}).Open()

	if errPool != nil {
		t.Errorf("Error starting pool: ", errPool)
		return
	}

	defer pool.Close()

	for i := 0; i < n_polls; i++ {
		if _, err := pool.SendWorkTimed(50, nil); err == nil {
		} else {
			t.Errorf("thread %v error: ", i, err)
		}
		outChan <- 1
	}

	for i := 0; i < n_polls; i++ {
		<-outChan
	}
}

func validateReturnInt(t *testing.T, expecting int, object interface{}) {
	if w, ok := object.(int); ok {
		if w != expecting {
			t.Errorf("Wrong, expected %v, got %v", expecting, w)
		}
	} else {
		t.Errorf("Wrong, expected int")
	}
}

func TestBasic(t *testing.T) {
	sizePool, repeats, sleepFor, margin := 16, 2, 250, 100
	outChan  := make(chan int, sizePool)

	runtime.GOMAXPROCS(runtime.NumCPU())

	pool, errPool := CreatePool(sizePool, func(object interface{}) interface{} {
		time.Sleep(time.Duration(sleepFor) * time.Millisecond)
		if w, ok := object.(int); ok {
			return w * 2
		}
		return "Not an int!"
	}).Open()

	if errPool != nil {
		t.Errorf("Error starting pool: ", errPool)
		return
	}

	defer pool.Close()

	for i := 0; i < sizePool * repeats; i++ {
		go func() {
			if out, err := pool.SendWork(50); err == nil {
				validateReturnInt (t, 100, out)
			} else {
				t.Errorf("Error returned: ", err)
			}
			outChan <- 1
		}()
	}

	before := time.Now()

	for i := 0; i < sizePool * repeats; i++ {
		<-outChan
	}

	taken    := float64( time.Since(before) ) / float64(time.Millisecond)
	expected := float64( sleepFor + margin ) * float64(repeats)

	if taken > expected {
		t.Errorf("Wrong, should have taken less than %v seconds, actually took %v", expected, taken)
	}
}

func TestGeneric(t *testing.T) {
	runtime.GOMAXPROCS(runtime.NumCPU())

	if pool, err := CreatePoolGeneric(10).Open(); err == nil {
		defer pool.Close()

		outChan  := make(chan int, 10)

		for i := 0; i < 10; i++ {
			go func(id int) {
				one, err := pool.SendWork(func() {
					outChan <- id
				})

				if err != nil {
					t.Errorf("Generic call timed out!")
				}

				if one != nil {
					if funcerr, ok := one.(error); ok {
						t.Errorf("Generic worker call: ", funcerr)
					} else {
						t.Errorf("Unexpected result from generic worker")
					}
				}
			}(i)
		}

		results := make([]int, 10)

		for i := 0; i < 10; i++ {
			value := <-outChan
			if results[value] != 0 || value > 9 || value < 0 {
				t.Errorf("duplicate or incorrect key: %v", value)
			}
			results[value] = 1
		}

	} else {
		t.Errorf("Error starting pool: ", err)
		return
	}

}

func TestExampleCase(t *testing.T) {
	outChan  := make(chan int, 10)
	runtime.GOMAXPROCS(runtime.NumCPU())

	pool, errPool := CreatePool(4, func(object interface{}) interface{} {
		if str, ok := object.(string); ok {
			return "job done: " + str
		}
		return nil
	}).Open()

	if errPool != nil {
		t.Errorf("Error starting pool: ", errPool)
		return
	}

	defer pool.Close()

	for i := 0; i < 10; i++ {
		go func() {
			if value, err := pool.SendWork("hello world"); err == nil {
				if _, ok := value.(string); ok {
				} else {
					t.Errorf("Not a string!")
				}
			} else {
				t.Errorf("Error returned: ", err)
			}
			outChan <- 1
		}()
	}

	for i := 0; i < 10; i++ {
		<-outChan
	}
}

type customWorker struct {
	jobsCompleted int
}

func (worker *customWorker) Ready() bool {
	return true
}

func (worker *customWorker) Job(data interface{}) interface{} {
	/* There's no need for thread safety paradigms here unless the data is being accessed from
	 * another goroutine outside of the pool.
	 */
	if outputStr, ok := data.(string); ok {
		(*worker).jobsCompleted++;
		return ("custom job done: " + outputStr)
	}
	return nil
}

func TestCustomWorkers(t *testing.T) {
	outChan  := make(chan int, 10)
	runtime.GOMAXPROCS(runtime.NumCPU())

	workers := make([]TunnyWorker, 4)
	for i, _ := range workers {
		workers[i] = &(customWorker{ jobsCompleted: 0 })
	}

	pool, errPool := CreateCustomPool(workers).Open()

	if errPool != nil {
		t.Errorf("Error starting pool: ", errPool)
		return
	}

	for i := 0; i < 10; i++ {
		/* Calling SendWork is thread safe, go ahead and call it from any goroutine.
		 * The call will block until a worker is ready and has completed the job.
		 */
		go func() {
			if value, err := pool.SendWork("hello world"); err == nil {
				if str, ok := value.(string); ok {
					if str != "custom job done: hello world" {
						t.Errorf("Unexpected output from custom worker")
					}
				} else {
					t.Errorf("Not a string!")
				}
			} else {
				t.Errorf("Error returned: ", err)
			}
			outChan <- 1
		}()
	}

	for i := 0; i < 10; i++ {
		<-outChan
	}

	/* After this call we should be able to guarantee that no other go routine is
	 * accessing the workers.
	 */
	pool.Close()

	totalJobs := 0
	for i := 0; i < len(workers); i++ {
		if custom, ok := workers[i].(*customWorker); ok {
			totalJobs += (*custom).jobsCompleted
		} else {
			t.Errorf("could not cast to customWorker")
		}
	}

	if totalJobs != 10 {
		t.Errorf("Total jobs expected: %v, actual: %v", 10, totalJobs)
	}
}

type customExtendedWorker struct {
	jobsCompleted int
	asleep        bool
}

func (worker *customExtendedWorker) Job(data interface{}) interface{} {
	if outputStr, ok := data.(string); ok {
		(*worker).jobsCompleted++;
		return ("custom job done: " + outputStr)
	}
	return nil
}

// Do 10 jobs and then stop.
func (worker *customExtendedWorker) Ready() bool {
	return !(*worker).asleep && ((*worker).jobsCompleted < 10)
}

func (worker *customExtendedWorker) Initialize() {
	(*worker).asleep = false
}

func (worker *customExtendedWorker) Terminate() {
	(*worker).asleep = true
}

func TestCustomExtendedWorkers(t *testing.T) {
	outChan  := make(chan int, 10)
	runtime.GOMAXPROCS(runtime.NumCPU())

	extWorkers   := make([]*customExtendedWorker, 4)
	tunnyWorkers := make([]TunnyWorker, 4)
	for i, _ := range tunnyWorkers {
		extWorkers  [i] = &(customExtendedWorker{ jobsCompleted: 0, asleep: true })
		tunnyWorkers[i] = extWorkers[i]
	}

	pool := CreateCustomPool(tunnyWorkers);

	for j := 0; j < 1; j++ {

		_, errPool := pool.Open()

		for i, _ := range extWorkers {
			if (*extWorkers[i]).asleep {
				t.Errorf("Worker is still asleep!")
			}
		}

		if errPool != nil {
			t.Errorf("Error starting pool: ", errPool)
			return
		}

		for i := 0; i < 40; i++ {
			/* Calling SendWork is thread safe, go ahead and call it from any goroutine.
			 * The call will block until a worker is ready and has completed the job.
			 */
			go func() {
				if value, err := pool.SendWork("hello world"); err == nil {
					if str, ok := value.(string); ok {
						if str != "custom job done: hello world" {
							t.Errorf("Unexpected output from custom worker")
						}
					} else {
						t.Errorf("Not a string!")
					}
				} else {
					t.Errorf("Error returned: ", err)
				}
				outChan <- 1
			}()
		}

		for i := 0; i < 40; i++ {
			<-outChan
		}

		/* After this call we should be able to guarantee that no other go routine is
		 * accessing the workers.
		 */
		pool.Close()

		expectedJobs := ((j + 1) * 10)
		for i, _ := range extWorkers {
			if (*extWorkers[i]).jobsCompleted != expectedJobs {
				t.Errorf( "Expected %v jobs completed, actually: %v",
					expectedJobs,
					(*extWorkers[i]).jobsCompleted,
				)
			}
			if !(*extWorkers[i]).asleep {
				t.Errorf("Worker is still awake!")
			}
		}
	}
}
<commit_msg>Tests now validate that goroutines aren't leaked<commit_after>/*
Copyright (c) 2014 Ashley Jeffs

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
*/

package tunny

import (
	"testing"
	"time"
	"runtime"
)

/* time.Sleep seems to be spawning a goroutine that doesn't stop. This is fine and stuff
 * but tunny tests do a comparisson of runtime.NumGoroutine before and after to make sure
 * that goroutines aren't being left open after .Close(), so this test simply opens that
 * goroutine beforehand.
 */
func TestTestTestyTestAllTheTestsNeedTesting(t * testing.T) {
	time.Sleep(time.Millisecond)
}

func TestTimeout(t *testing.T) {
	outChan  := make(chan int, 3)
	routines := runtime.NumGoroutine()

	pool, errPool := CreatePool(1, func(object interface{}) interface{} {
		time.Sleep(500 * time.Millisecond)
		return nil
	}).Open()

	if errPool != nil {
		t.Errorf("Error starting pool: ", errPool)
		return
	}

	before := time.Now()

	go func() {
		if _, err := pool.SendWorkTimed(200, nil); err == nil {
			t.Errorf("No timeout triggered thread one")
		} else {
			taken := ( time.Since(before) / time.Millisecond )
			if taken > 210 {
				t.Errorf("Time taken at thread one: ", taken, ", with error: ", err)
			}
		}
		outChan <- 1

		go func() {
			if _, err := pool.SendWork(nil); err == nil {
			} else {
				t.Errorf("Error at thread three: ", err)
			}
			outChan <- 1
		}()
	}()

	go func() {
		if _, err := pool.SendWorkTimed(200, nil); err == nil {
			t.Errorf("No timeout triggered thread two")
		} else {
			taken := ( time.Since(before) / time.Millisecond )
			if taken > 210 {
				t.Errorf("Time taken at thread two: ", taken, ", with error: ", err)
			}
		}
		outChan <- 1
	}()

	for i := 0; i < 3; i++ {
		<-outChan
	}

	pool.Close()

	if routines != runtime.NumGoroutine() {
		t.Errorf("Excess goroutines: %v", runtime.NumGoroutine() - routines)
	}
}

func TestTimeoutRequests(t *testing.T) {
	n_polls  := 200
	outChan  := make(chan int, n_polls)
	routines := runtime.NumGoroutine()

	pool, errPool := CreatePool(1, func(object interface{}) interface{} {
		time.Sleep(time.Millisecond)
		return nil
	}).Open()

	if errPool != nil {
		t.Errorf("Error starting pool: ", errPool)
		return
	}

	for i := 0; i < n_polls; i++ {
		if _, err := pool.SendWorkTimed(50, nil); err == nil {
		} else {
			t.Errorf("thread %v error: ", i, err)
		}
		outChan <- 1
	}

	for i := 0; i < n_polls; i++ {
		<-outChan
	}

	pool.Close()

	if routines != runtime.NumGoroutine() {
		t.Errorf("Excess goroutines: %v", runtime.NumGoroutine() - routines)
	}
}

func validateReturnInt(t *testing.T, expecting int, object interface{}) {
	if w, ok := object.(int); ok {
		if w != expecting {
			t.Errorf("Wrong, expected %v, got %v", expecting, w)
		}
	} else {
		t.Errorf("Wrong, expected int")
	}
}

func TestBasic(t *testing.T) {
	sizePool, repeats, sleepFor, margin := 16, 2, 250, 100
	outChan  := make(chan int, sizePool)
	routines := runtime.NumGoroutine()

	runtime.GOMAXPROCS(runtime.NumCPU())

	pool, errPool := CreatePool(sizePool, func(object interface{}) interface{} {
		time.Sleep(time.Duration(sleepFor) * time.Millisecond)
		if w, ok := object.(int); ok {
			return w * 2
		}
		return "Not an int!"
	}).Open()

	if errPool != nil {
		t.Errorf("Error starting pool: ", errPool)
		return
	}

	for i := 0; i < sizePool * repeats; i++ {
		go func() {
			if out, err := pool.SendWork(50); err == nil {
				validateReturnInt (t, 100, out)
			} else {
				t.Errorf("Error returned: ", err)
			}
			outChan <- 1
		}()
	}

	before := time.Now()

	for i := 0; i < sizePool * repeats; i++ {
		<-outChan
	}

	taken    := float64( time.Since(before) ) / float64(time.Millisecond)
	expected := float64( sleepFor + margin ) * float64(repeats)

	if taken > expected {
		t.Errorf("Wrong, should have taken less than %v seconds, actually took %v", expected, taken)
	}

	pool.Close()

	if routines != runtime.NumGoroutine() {
		t.Errorf("Excess goroutines: %v", runtime.NumGoroutine() - routines)
	}
}

func TestGeneric(t *testing.T) {
	routines := runtime.NumGoroutine()
	runtime.GOMAXPROCS(runtime.NumCPU())

	if pool, err := CreatePoolGeneric(10).Open(); err == nil {
		outChan  := make(chan int, 10)

		for i := 0; i < 10; i++ {
			go func(id int) {
				one, err := pool.SendWork(func() {
					outChan <- id
				})

				if err != nil {
					t.Errorf("Generic call timed out!")
				}

				if one != nil {
					if funcerr, ok := one.(error); ok {
						t.Errorf("Generic worker call: ", funcerr)
					} else {
						t.Errorf("Unexpected result from generic worker")
					}
				}
			}(i)
		}

		results := make([]int, 10)

		for i := 0; i < 10; i++ {
			value := <-outChan
			if results[value] != 0 || value > 9 || value < 0 {
				t.Errorf("duplicate or incorrect key: %v", value)
			}
			results[value] = 1
		}

		pool.Close()

	} else {
		t.Errorf("Error starting pool: ", err)
		return
	}

	if routines != runtime.NumGoroutine() {
		t.Errorf("Excess goroutines: %v", runtime.NumGoroutine() - routines)
	}
}

func TestExampleCase(t *testing.T) {
	outChan  := make(chan int, 10)
	routines := runtime.NumGoroutine()
	runtime.GOMAXPROCS(runtime.NumCPU())

	pool, errPool := CreatePool(4, func(object interface{}) interface{} {
		if str, ok := object.(string); ok {
			return "job done: " + str
		}
		return nil
	}).Open()

	if errPool != nil {
		t.Errorf("Error starting pool: ", errPool)
		return
	}

	for i := 0; i < 10; i++ {
		go func() {
			if value, err := pool.SendWork("hello world"); err == nil {
				if _, ok := value.(string); ok {
				} else {
					t.Errorf("Not a string!")
				}
			} else {
				t.Errorf("Error returned: ", err)
			}
			outChan <- 1
		}()
	}

	for i := 0; i < 10; i++ {
		<-outChan
	}

	pool.Close()

	if routines != runtime.NumGoroutine() {
		t.Errorf("Excess goroutines: %v", runtime.NumGoroutine() - routines)
	}
}

type customWorker struct {
	jobsCompleted int
}

func (worker *customWorker) Ready() bool {
	return true
}

func (worker *customWorker) Job(data interface{}) interface{} {
	/* There's no need for thread safety paradigms here unless the data is being accessed from
	 * another goroutine outside of the pool.
	 */
	if outputStr, ok := data.(string); ok {
		(*worker).jobsCompleted++;
		return ("custom job done: " + outputStr)
	}
	return nil
}

func TestCustomWorkers(t *testing.T) {
	outChan  := make(chan int, 10)
	routines := runtime.NumGoroutine()

	runtime.GOMAXPROCS(runtime.NumCPU())

	workers := make([]TunnyWorker, 4)
	for i, _ := range workers {
		workers[i] = &(customWorker{ jobsCompleted: 0 })
	}

	pool, errPool := CreateCustomPool(workers).Open()

	if errPool != nil {
		t.Errorf("Error starting pool: ", errPool)
		return
	}

	for i := 0; i < 10; i++ {
		/* Calling SendWork is thread safe, go ahead and call it from any goroutine.
		 * The call will block until a worker is ready and has completed the job.
		 */
		go func() {
			if value, err := pool.SendWork("hello world"); err == nil {
				if str, ok := value.(string); ok {
					if str != "custom job done: hello world" {
						t.Errorf("Unexpected output from custom worker")
					}
				} else {
					t.Errorf("Not a string!")
				}
			} else {
				t.Errorf("Error returned: ", err)
			}
			outChan <- 1
		}()
	}

	for i := 0; i < 10; i++ {
		<-outChan
	}

	/* After this call we should be able to guarantee that no other go routine is
	 * accessing the workers.
	 */
	pool.Close()

	totalJobs := 0
	for i := 0; i < len(workers); i++ {
		if custom, ok := workers[i].(*customWorker); ok {
			totalJobs += (*custom).jobsCompleted
		} else {
			t.Errorf("could not cast to customWorker")
		}
	}

	if totalJobs != 10 {
		t.Errorf("Total jobs expected: %v, actual: %v", 10, totalJobs)
	}

	if routines != runtime.NumGoroutine() {
		t.Errorf("Excess goroutines: %v", runtime.NumGoroutine() - routines)
	}
}

type customExtendedWorker struct {
	jobsCompleted int
	asleep        bool
}

func (worker *customExtendedWorker) Job(data interface{}) interface{} {
	if outputStr, ok := data.(string); ok {
		(*worker).jobsCompleted++;
		return ("custom job done: " + outputStr)
	}
	return nil
}

// Do 10 jobs and then stop.
func (worker *customExtendedWorker) Ready() bool {
	return !(*worker).asleep && ((*worker).jobsCompleted < 10)
}

func (worker *customExtendedWorker) Initialize() {
	(*worker).asleep = false
}

func (worker *customExtendedWorker) Terminate() {
	(*worker).asleep = true
}

func TestCustomExtendedWorkers(t *testing.T) {
	outChan  := make(chan int, 10)
	routines := runtime.NumGoroutine()
	runtime.GOMAXPROCS(runtime.NumCPU())

	extWorkers   := make([]*customExtendedWorker, 4)
	tunnyWorkers := make([]TunnyWorker, 4)
	for i, _ := range tunnyWorkers {
		extWorkers  [i] = &(customExtendedWorker{ jobsCompleted: 0, asleep: true })
		tunnyWorkers[i] = extWorkers[i]
	}

	pool := CreateCustomPool(tunnyWorkers);

	for j := 0; j < 1; j++ {

		_, errPool := pool.Open()

		for i, _ := range extWorkers {
			if (*extWorkers[i]).asleep {
				t.Errorf("Worker is still asleep!")
			}
		}

		if errPool != nil {
			t.Errorf("Error starting pool: ", errPool)
			return
		}

		for i := 0; i < 40; i++ {
			/* Calling SendWork is thread safe, go ahead and call it from any goroutine.
			 * The call will block until a worker is ready and has completed the job.
			 */
			go func() {
				if value, err := pool.SendWork("hello world"); err == nil {
					if str, ok := value.(string); ok {
						if str != "custom job done: hello world" {
							t.Errorf("Unexpected output from custom worker")
						}
					} else {
						t.Errorf("Not a string!")
					}
				} else {
					t.Errorf("Error returned: ", err)
				}
				outChan <- 1
			}()
		}

		for i := 0; i < 40; i++ {
			<-outChan
		}

		/* After this call we should be able to guarantee that no other go routine is
		 * accessing the workers.
		 */
		pool.Close()

		expectedJobs := ((j + 1) * 10)
		for i, _ := range extWorkers {
			if (*extWorkers[i]).jobsCompleted != expectedJobs {
				t.Errorf( "Expected %v jobs completed, actually: %v",
					expectedJobs,
					(*extWorkers[i]).jobsCompleted,
				)
			}
			if !(*extWorkers[i]).asleep {
				t.Errorf("Worker is still awake!")
			}
		}
	}

	if routines != runtime.NumGoroutine() {
		t.Errorf("Excess goroutines: %v", runtime.NumGoroutine() - routines)
	}
}
<|endoftext|>


<commit_before>// clock counts down to or up from a target time.
package main

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

func main() {
	target := time.Date(2016, 8, 8, 0, 0, 0, 0, time.Local)
	motto := "Just Go"
	printTargetTime(target, motto)
	exitOnEnterKey()

	var previous time.Time
	for {
		now := time.Now()
		now = now.Add(time.Duration(-now.Nanosecond())) // truncate to second
		if now != previous {
			previous = now
			remaining := target.Sub(now)
			printTimeRemaining(now, remaining)
		}
		time.Sleep(50 * time.Millisecond)
	}
}

func exitOnEnterKey() {
	go func() {
		buf := make([]byte, 1)
		_, _ = os.Stdin.Read(buf)
		os.Exit(0)
	}()
}

const (
	indent         = "\t"
	highlightStart = "\x1b[1;35m"
	highlightEnd   = "\x1b[0m"
)

func printTargetTime(target time.Time, motto string) {
	fmt.Print(indent, highlightStart, motto, highlightEnd, "\n")
	fmt.Print(indent, target.Format(time.UnixDate), "\n")
}

func printTimeRemaining(now time.Time, remaining time.Duration) {
	var sign string
	if remaining > 0 {
		sign = "-" // countdown is "T minus..."
	} else {
		sign = "+" // count up is "T plus..."
		remaining = -remaining
	}

	var days int
	if remaining >= 24*time.Hour {
		days = int(remaining / (24 * time.Hour))
		remaining = remaining % (24 * time.Hour)
	}

	fmt.Print(indent, now.Format(time.UnixDate), "  ", sign)
	if days > 0 {
		fmt.Print(days, "d")
	}
	fmt.Print(remaining, "          \r")
	os.Stdout.Sync()
}
<commit_msg>Speculative Go 1.7<commit_after>// clock counts down to or up from a target time.
package main

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

func main() {
	target := time.Date(2016, 8, 15, 0, 0, 0, 0, time.Local)
	motto := "Just Go"
	printTargetTime(target, motto)
	exitOnEnterKey()

	var previous time.Time
	for {
		now := time.Now()
		now = now.Add(time.Duration(-now.Nanosecond())) // truncate to second
		if now != previous {
			previous = now
			remaining := target.Sub(now)
			printTimeRemaining(now, remaining)
		}
		time.Sleep(50 * time.Millisecond)
	}
}

func exitOnEnterKey() {
	go func() {
		buf := make([]byte, 1)
		_, _ = os.Stdin.Read(buf)
		os.Exit(0)
	}()
}

const (
	indent         = "\t"
	highlightStart = "\x1b[1;35m"
	highlightEnd   = "\x1b[0m"
)

func printTargetTime(target time.Time, motto string) {
	fmt.Print(indent, highlightStart, motto, highlightEnd, "\n")
	fmt.Print(indent, target.Format(time.UnixDate), "\n")
}

func printTimeRemaining(now time.Time, remaining time.Duration) {
	var sign string
	if remaining > 0 {
		sign = "-" // countdown is "T minus..."
	} else {
		sign = "+" // count up is "T plus..."
		remaining = -remaining
	}

	var days int
	if remaining >= 24*time.Hour {
		days = int(remaining / (24 * time.Hour))
		remaining = remaining % (24 * time.Hour)
	}

	fmt.Print(indent, now.Format(time.UnixDate), "  ", sign)
	if days > 0 {
		fmt.Print(days, "d")
	}
	fmt.Print(remaining, "          \r")
	os.Stdout.Sync()
}
<|endoftext|>


<commit_before>package main

import (
//    "fmt"
)

type UDP_Read_Manager struct {
	ipAddress string
	reader    *IP_Reader
	buff      map[uint16](chan []byte)
}

type UDP_Reader struct {
	manager *UDP_Read_Manager
	bytes   <-chan []byte
	port    uint16 // ports
}

func NewUDP_Read_Manager(ip string) (*UDP_Read_Manager, error) {
	nr, err := NewNetwork_Reader()
	if err != nil {
		return nil, err
	}

	ipr, err := nr.NewIP_Reader(ip, 17) // 17 for UDP
	if err != nil {
		return nil, err
	}

	x := &UDP_Read_Manager{
		reader:    ipr,
		buff:      make(map[uint16](chan []byte)),
		ipAddress: ip,
	}

	go x.readAll()

	return x, nil
}

func (x *UDP_Read_Manager) readAll() {
	for {
		_, payload, err := x.reader.ReadFrom()
		if err != nil {
			continue
		}
		//fmt.Println(b)
		//fmt.Println("UDP header and payload: ", payload)

		dest := (((uint16)(payload[2])) * 256) + ((uint16)(payload[3]))
		//fmt.Println(dest)
		//fmt.Println(payload)
		//
		//fmt.Println(x.buff)
		c, ok := x.buff[dest]
		//fmt.Println(ok)
		payload = payload[8:]
		if ok {
			go func() { c <- payload }()
		} else {
			// drop packet
		}
	}
}

func (x *UDP_Read_Manager) NewUDP(port uint16) (*UDP_Reader, error) {
	x.buff[port] = make(chan []byte)
	return &UDP_Reader{port: port, bytes: x.buff[port], manager: x}, nil
}

func (c *UDP_Reader) read(size int) ([]byte, error) {
	data := <-c.bytes
	if len(data) > size {
		data = data[:size]
	}
	return data, nil
}

func (c *UDP_Reader) close() error {
	delete(c.manager.buff, c.port)
	return nil
}
<commit_msg>automatic commit<commit_after>package main

import (
//    "fmt"
)

const MAX_UDP_PACKET_LEN = 65507

type UDP_Read_Manager struct {
	ipAddress string
	reader    *IP_Reader
	buff      map[uint16](chan []byte)
}

type UDP_Reader struct {
	manager *UDP_Read_Manager
	bytes   <-chan []byte
	port    uint16 // ports
}

func NewUDP_Read_Manager(ip string) (*UDP_Read_Manager, error) {
	nr, err := NewNetwork_Reader()
	if err != nil {
		return nil, err
	}

	ipr, err := nr.NewIP_Reader(ip, 17) // 17 for UDP
	if err != nil {
		return nil, err
	}

	x := &UDP_Read_Manager{
		reader:    ipr,
		buff:      make(map[uint16](chan []byte)),
		ipAddress: ip,
	}

	go x.readAll()

	return x, nil
}

func (x *UDP_Read_Manager) readAll() {
	for {
		_, payload, err := x.reader.ReadFrom()
		if err != nil {
			continue
		}
		//fmt.Println(b)
		//fmt.Println("UDP header and payload: ", payload)

		dest := (((uint16)(payload[2])) * 256) + ((uint16)(payload[3]))
		//fmt.Println(dest)
		//fmt.Println(payload)
		//
		//fmt.Println(x.buff)
		c, ok := x.buff[dest]
		//fmt.Println(ok)
		payload = payload[8:]
		if ok {
			go func() { c <- payload }()
		} else {
			// drop packet
		}
	}
}

func (x *UDP_Read_Manager) NewUDP(port uint16) (*UDP_Reader, error) {
	x.buff[port] = make(chan []byte)
	return &UDP_Reader{port: port, bytes: x.buff[port], manager: x}, nil
}

func (c *UDP_Reader) read(size int) ([]byte, error) {
	data := <-c.bytes
	if len(data) > size {
		data = data[:size]
	}
	return data, nil
}

func (c *UDP_Reader) close() error {
	delete(c.manager.buff, c.port)
	return nil
}
<|endoftext|>


<commit_before>// Test GoAWK Lexer

package lexer_test

import (
	"fmt"
	"strconv"
	"strings"
	"testing"

	. "github.com/benhoyt/goawk/lexer"
)

// TODO: add other lexer tests

func TestNumber(t *testing.T) {
	tests := []struct {
		input  string
		output string
	}{
		{"0", "1:1 number 0"},
		{"9", "1:1 number 9"},
		{" 0 ", "1:2 number 0"},
		{"\n  1", "1:1 newline , 2:3 number 1"},
		{"1234", "1:1 number 1234"},
		{".5", "1:1 number .5"},
		{".5e1", "1:1 number .5e1"},
		{"5e+1", "1:1 number 5e+1"},
		{"5e-1", "1:1 number 5e-1"},
		{"0.", "1:1 number 0."},
		{"42e", "1:1 number 42e"},
		{"4.2e", "1:1 number 4.2e"},
		{"1.e3", "1:1 number 1.e3"},
		{"1.e3", "1:1 number 1.e3"},
		{"1e3foo", "1:1 number 1e3, 1:4 name foo"},
		{"1e3+", "1:1 number 1e3, 1:4 + "},
		{"1e3.4", "1:1 number 1e3, 1:4 number .4"},
		{"42@", "1:1 number 42, 1:3 <illegal> unexpected '@'"},
		{"0..", "1:1 number 0., 1:4 <illegal> expected digits"},
		{".", "1:2 <illegal> expected digits"},
	}
	for _, test := range tests {
		t.Run(test.input, func(t *testing.T) {
			l := NewLexer([]byte(test.input))
			strs := []string{}
			for {
				pos, tok, val := l.Scan()
				if tok == EOF {
					break
				}
				if tok == NUMBER {
					// Ensure ParseFloat() works, as that's what our
					// parser uses to convert
					trimmed := strings.TrimRight(val, "eE")
					_, err := strconv.ParseFloat(trimmed, 64)
					if err != nil {
						t.Fatalf("couldn't parse float: %q", val)
					}
				}
				strs = append(strs, fmt.Sprintf("%d:%d %s %s", pos.Line, pos.Column, tok, val))
			}
			output := strings.Join(strs, ", ")
			if output != test.output {
				t.Errorf("expected %q, got %q", test.output, output)
			}
		})
	}
}

func TestStringMethod(t *testing.T) {
	input := "# comment line\n" +
		"+ += && = : , -- /\n/= $ == >= > ++ { [ < ( #\n" +
		"<= ~ % %= * *= !~ ! != || ^ ^= ** **= ? } ] ) ; - -= " +
		"BEGIN break continue delete do else END exit " +
		"for if in next print printf return while " +
		"atan2 cos exp gsub index int length log match rand " +
		"sin split sprintf sqrt srand sub substr tolower toupper " +
		"x \"str\\n\" 1234\n" +
		"/foo/\n" +
		"@ ."

	strs := make([]string, 0, LAST+1)
	seen := make([]bool, LAST+1)
	l := NewLexer([]byte(input))
	for {
		_, tok, _ := l.Scan()
		strs = append(strs, tok.String())
		seen[int(tok)] = true
		if tok == EOF {
			break
		}
	}
	output := strings.Join(strs, " ")

	expected := "newline " +
		"+ += && = : , -- / newline /= $ == >= > ++ { [ < ( newline " +
		"<= ~ % %= * *= !~ ! != || ^ ^= ^ ^= ? } ] ) ; - -= " +
		"BEGIN break continue delete do else END exit " +
		"for if in next print printf return while " +
		"atan2 cos exp gsub index int length log match rand " +
		"sin split sprintf sqrt srand sub substr tolower toupper " +
		"name string number newline " +
		"regex newline " +
		"<illegal> <illegal> EOF"
	if output != expected {
		t.Errorf("expected %q, got %q", expected, output)
	}

	var unsupportedTokens = map[Token]bool{
		F_CLOSE:  true,
		F_SYSTEM: true,
		FUNCTION: true,
		GETLINE:  true,
	}
	for i, s := range seen {
		if !unsupportedTokens[Token(i)] && !s && Token(i) != CONCAT {
			t.Errorf("token %s (%d) not seen", Token(i), i)
		}
	}
}
<commit_msg>Add unsupported tokens to test string so we don't forget them later<commit_after>// Test GoAWK Lexer

package lexer_test

import (
	"fmt"
	"strconv"
	"strings"
	"testing"

	. "github.com/benhoyt/goawk/lexer"
)

// TODO: add other lexer tests

func TestNumber(t *testing.T) {
	tests := []struct {
		input  string
		output string
	}{
		{"0", "1:1 number 0"},
		{"9", "1:1 number 9"},
		{" 0 ", "1:2 number 0"},
		{"\n  1", "1:1 newline , 2:3 number 1"},
		{"1234", "1:1 number 1234"},
		{".5", "1:1 number .5"},
		{".5e1", "1:1 number .5e1"},
		{"5e+1", "1:1 number 5e+1"},
		{"5e-1", "1:1 number 5e-1"},
		{"0.", "1:1 number 0."},
		{"42e", "1:1 number 42e"},
		{"4.2e", "1:1 number 4.2e"},
		{"1.e3", "1:1 number 1.e3"},
		{"1.e3", "1:1 number 1.e3"},
		{"1e3foo", "1:1 number 1e3, 1:4 name foo"},
		{"1e3+", "1:1 number 1e3, 1:4 + "},
		{"1e3.4", "1:1 number 1e3, 1:4 number .4"},
		{"42@", "1:1 number 42, 1:3 <illegal> unexpected '@'"},
		{"0..", "1:1 number 0., 1:4 <illegal> expected digits"},
		{".", "1:2 <illegal> expected digits"},
	}
	for _, test := range tests {
		t.Run(test.input, func(t *testing.T) {
			l := NewLexer([]byte(test.input))
			strs := []string{}
			for {
				pos, tok, val := l.Scan()
				if tok == EOF {
					break
				}
				if tok == NUMBER {
					// Ensure ParseFloat() works, as that's what our
					// parser uses to convert
					trimmed := strings.TrimRight(val, "eE")
					_, err := strconv.ParseFloat(trimmed, 64)
					if err != nil {
						t.Fatalf("couldn't parse float: %q", val)
					}
				}
				strs = append(strs, fmt.Sprintf("%d:%d %s %s", pos.Line, pos.Column, tok, val))
			}
			output := strings.Join(strs, ", ")
			if output != test.output {
				t.Errorf("expected %q, got %q", test.output, output)
			}
		})
	}
}

func TestStringMethod(t *testing.T) {
	input := "# comment line\n" +
		"+ += && = : , -- /\n/= $ == >= > ++ { [ < ( #\n" +
		"<= ~ % %= * *= !~ ! != || ^ ^= ** **= ? } ] ) ; - -= " +
		"BEGIN break continue delete do else END exit " +
		"for function getline if in next print printf return while " +
		"atan2 close cos exp gsub index int length log match rand " +
		"sin split sprintf sqrt srand sub substr system tolower toupper " +
		"x \"str\\n\" 1234\n" +
		"/foo/\n" +
		"@ ."

	strs := make([]string, 0, LAST+1)
	seen := make([]bool, LAST+1)
	l := NewLexer([]byte(input))
	for {
		_, tok, _ := l.Scan()
		strs = append(strs, tok.String())
		seen[int(tok)] = true
		if tok == EOF {
			break
		}
	}
	output := strings.Join(strs, " ")

	expected := "newline " +
		"+ += && = : , -- / newline /= $ == >= > ++ { [ < ( newline " +
		"<= ~ % %= * *= !~ ! != || ^ ^= ^ ^= ? } ] ) ; - -= " +
		"BEGIN break continue delete do else END exit " +
		"for <illegal> <illegal> if in next print printf return while " +
		"atan2 <illegal> cos exp gsub index int length log match rand " +
		"sin split sprintf sqrt srand sub substr <illegal> tolower toupper " +
		"name string number newline " +
		"regex newline " +
		"<illegal> <illegal> EOF"
	if output != expected {
		t.Errorf("expected %q, got %q", expected, output)
	}

	var unsupportedTokens = map[Token]bool{
		F_CLOSE:  true,
		F_SYSTEM: true,
		FUNCTION: true,
		GETLINE:  true,
	}
	for i, s := range seen {
		if !unsupportedTokens[Token(i)] && !s && Token(i) != CONCAT {
			t.Errorf("token %s (%d) not seen", Token(i), i)
		}
	}
}
<|endoftext|>


<commit_before>package golanglrutimedcache

import (
	"strings"
	"sync"
	"time"

	"github.com/hashicorp/golang-lru"
)

type cacheElement struct {
	loadTime int64  //The time when the element was loaded from database
	keyName  string //The name of the key (if known)
	element  interface{}
}

//TimedCache is an lru cache with an additional expiration time for the data in each element
//It is not a very clever cache, and does not factor in time expiration in lru
type TimedCache struct {
	sync.Mutex //Necessary to use the nameMap
	c          *lru.Cache
	expireTime int64
	nameMap    map[string]int64 //A map that goes from name to ID
}

//NewTimedCache returns the TimedCache with given parameters
func NewTimedCache(size int, expireTime int64, err error) (*TimedCache, error) {
	if err != nil {
		return nil, err
	}

	tc := &TimedCache{sync.Mutex{}, nil, expireTime, make(map[string]int64)}

	//Add an eviction
	tc.c, err = lru.NewWithEvict(size, func(key interface{}, value interface{}) {
		tc.Lock()
		delete(tc.nameMap, value.(cacheElement).keyName)
		tc.Unlock()
	})
	return tc, err
}

//Set adds a new element to the cache
func (tc *TimedCache) Set(name string, id int64, value interface{}) {
	if _, keyName, ok := tc.GetByID(id); ok && keyName != "" && keyName != name {
		tc.UnlinkName(keyName)
	}
	tc.c.Add(id, cacheElement{time.Now().Unix(), name, value})
	tc.Lock()
	tc.nameMap[name] = id
	tc.Unlock()
}

//Update an element without changing the associated name
func (tc *TimedCache) Update(id int64, value interface{}) {
	_, keyName, _ := tc.GetByID(id)
	tc.c.Add(id, cacheElement{time.Now().Unix(), keyName, value})
}

//SetID sets a new element without knowing its name
func (tc *TimedCache) SetID(id int64, value interface{}) {
	if _, keyName, ok := tc.GetByID(id); ok && keyName != "" {
		tc.UnlinkName(keyName)
	}
	tc.c.Add(id, cacheElement{time.Now().Unix(), "", value})
}

//UnlinkName unlinks the name fro mthe ID
func (tc *TimedCache) UnlinkName(keyName string) {
	tc.Lock()
	val, ok := tc.nameMap[keyName]
	delete(tc.nameMap, keyName)
	tc.Unlock()
	if ok {
		v, ok := tc.c.Get(val)
		if ok {
			ce := v.(cacheElement)
			ce.keyName = ""
			tc.c.Add(val, ce)
		}

	}
}

//GetNameID gets the ID associated with a name
func (tc *TimedCache) GetNameID(keyName string) (nameID int64, ok bool) {
	tc.Lock()
	defer tc.Unlock()
	val, ok := tc.nameMap[keyName]
	return val, ok
}

//UnlinkNamePrefix removes all names with the given prefix. This is not a very efficient way to do it,
//since it loops through the entire cache unlinking names.  It locks queries by name for the entire time.
//it might be a better idea to just let names expire, since if a user is deleted, any database operation on the user
//will fail miserably. Writing to a stream belonging to a deleted user is still possible that way, but
//Redis has a cleaning process that runs periodically.
func (tc *TimedCache) UnlinkNamePrefix(namePrefix string) {
	tc.Lock()
	for key := range tc.nameMap {
		if strings.HasPrefix(key, namePrefix) {
			delete(tc.nameMap, key)
		}
	}
	tc.Unlock()
}

//GetByID gets an element if it is cached and not yet expired
func (tc *TimedCache) GetByID(key int64) (value interface{}, keyName string, ok bool) {
	v, ok := tc.c.Get(key)
	if !ok {
		return nil, "", false
	}
	ce := v.(cacheElement)
	if ce.loadTime+tc.expireTime <= time.Now().Unix() {
		//The element is expired - it shouldn't be in the cache
		tc.UnlinkName(ce.keyName)
		tc.c.Remove(key)
		return nil, "", false
	}
	return ce.element, ce.keyName, true
}

//GetByName gets an element if it is in the cache and not expired
func (tc *TimedCache) GetByName(name string) (value interface{}, ok bool) {
	key, ok := tc.GetNameID(name)
	if !ok {
		return nil, false
	}
	value, keyName, ok := tc.GetByID(key)
	if keyName != name { //If ok is false, keyName = ""
		tc.UnlinkName(name) //Make sure that this is a matching name and ID
		ok = false
	}
	return value, ok
}

//RemoveID the given key from the cache by ID
func (tc *TimedCache) RemoveID(key int64) {
	v, ok := tc.c.Get(key)
	if ok {
		tc.UnlinkName(v.(cacheElement).keyName)
	}
	tc.c.Remove(key)
}

//RemoveName deletes the cache element by its name
func (tc *TimedCache) RemoveName(name string) {
	id, ok := tc.GetNameID(name)
	tc.UnlinkName(name)
	if ok {
		tc.RemoveID(id)
	}
}

//PurgeNames clears all the names from the cache, but leaves the cache elements accessible by ID
func (tc *TimedCache) PurgeNames() {
	tc.Lock()
	tc.nameMap = make(map[string]int64)
	tc.Unlock()
}

//Purge clears all elements in the cache
func (tc *TimedCache) Purge() {
	tc.PurgeNames()
	tc.c.Purge()
}
<commit_msg>Removed references to original source<commit_after>package golanglrutimedcache

import (
	"strings"
	"sync"
	"time"

	"github.com/hashicorp/golang-lru"
)

type cacheElement struct {
	loadTime int64  //The time when the element was loaded from database
	keyName  string //The name of the key (if known)
	element  interface{}
}

//TimedCache is an lru cache with an additional expiration time for the data in each element
//It is not a very clever cache, and does not factor in time expiration in lru
type TimedCache struct {
	sync.Mutex //Necessary to use the nameMap
	c          *lru.Cache
	expireTime int64
	nameMap    map[string]int64 //A map that goes from name to ID
}

//NewTimedCache returns the TimedCache with given parameters
func NewTimedCache(size int, expireTime int64, err error) (*TimedCache, error) {
	if err != nil {
		return nil, err
	}

	tc := &TimedCache{sync.Mutex{}, nil, expireTime, make(map[string]int64)}

	//Add an eviction
	tc.c, err = lru.NewWithEvict(size, func(key interface{}, value interface{}) {
		tc.Lock()
		delete(tc.nameMap, value.(cacheElement).keyName)
		tc.Unlock()
	})
	return tc, err
}

//Set adds a new element to the cache
func (tc *TimedCache) Set(name string, id int64, value interface{}) {
	if _, keyName, ok := tc.GetByID(id); ok && keyName != "" && keyName != name {
		tc.UnlinkName(keyName)
	}
	tc.c.Add(id, cacheElement{time.Now().Unix(), name, value})
	tc.Lock()
	tc.nameMap[name] = id
	tc.Unlock()
}

//Update an element without changing the associated name
func (tc *TimedCache) Update(id int64, value interface{}) {
	_, keyName, _ := tc.GetByID(id)
	tc.c.Add(id, cacheElement{time.Now().Unix(), keyName, value})
}

//SetID sets a new element without knowing its name
func (tc *TimedCache) SetID(id int64, value interface{}) {
	if _, keyName, ok := tc.GetByID(id); ok && keyName != "" {
		tc.UnlinkName(keyName)
	}
	tc.c.Add(id, cacheElement{time.Now().Unix(), "", value})
}

//UnlinkName unlinks the name fro mthe ID
func (tc *TimedCache) UnlinkName(keyName string) {
	tc.Lock()
	val, ok := tc.nameMap[keyName]
	delete(tc.nameMap, keyName)
	tc.Unlock()
	if ok {
		v, ok := tc.c.Get(val)
		if ok {
			ce := v.(cacheElement)
			ce.keyName = ""
			tc.c.Add(val, ce)
		}

	}
}

//GetNameID gets the ID associated with a name
func (tc *TimedCache) GetNameID(keyName string) (nameID int64, ok bool) {
	tc.Lock()
	defer tc.Unlock()
	val, ok := tc.nameMap[keyName]
	return val, ok
}

//UnlinkNamePrefix removes all names with the given prefix. This is not a very efficient way to do it,
//since it loops through the entire cache unlinking names.  It locks queries by name for the entire time.
func (tc *TimedCache) UnlinkNamePrefix(namePrefix string) {
	tc.Lock()
	for key := range tc.nameMap {
		if strings.HasPrefix(key, namePrefix) {
			delete(tc.nameMap, key)
		}
	}
	tc.Unlock()
}

//GetByID gets an element if it is cached and not yet expired
func (tc *TimedCache) GetByID(key int64) (value interface{}, keyName string, ok bool) {
	v, ok := tc.c.Get(key)
	if !ok {
		return nil, "", false
	}
	ce := v.(cacheElement)
	if ce.loadTime+tc.expireTime <= time.Now().Unix() {
		//The element is expired - it shouldn't be in the cache
		tc.UnlinkName(ce.keyName)
		tc.c.Remove(key)
		return nil, "", false
	}
	return ce.element, ce.keyName, true
}

//GetByName gets an element if it is in the cache and not expired
func (tc *TimedCache) GetByName(name string) (value interface{}, ok bool) {
	key, ok := tc.GetNameID(name)
	if !ok {
		return nil, false
	}
	value, keyName, ok := tc.GetByID(key)
	if keyName != name { //If ok is false, keyName = ""
		tc.UnlinkName(name) //Make sure that this is a matching name and ID
		ok = false
	}
	return value, ok
}

//RemoveID the given key from the cache by ID
func (tc *TimedCache) RemoveID(key int64) {
	v, ok := tc.c.Get(key)
	if ok {
		tc.UnlinkName(v.(cacheElement).keyName)
	}
	tc.c.Remove(key)
}

//RemoveName deletes the cache element by its name
func (tc *TimedCache) RemoveName(name string) {
	id, ok := tc.GetNameID(name)
	tc.UnlinkName(name)
	if ok {
		tc.RemoveID(id)
	}
}

//PurgeNames clears all the names from the cache, but leaves the cache elements accessible by ID
func (tc *TimedCache) PurgeNames() {
	tc.Lock()
	tc.nameMap = make(map[string]int64)
	tc.Unlock()
}

//Purge clears all elements in the cache
func (tc *TimedCache) Purge() {
	tc.PurgeNames()
	tc.c.Purge()
}
<|endoftext|>


<commit_before>package cache

import (
	"fmt"
	"time"
	"net/http"
	"net/http/httptest"
	"gopkg.in/redis.v4"
	"gopkg.in/vmihailenco/msgpack.v2"
	"github.com/mholt/caddy/caddyhttp/httpserver"
)

type CacheHandler struct {
	Client *redis.Client
	Next   httpserver.Handler
}

type CachedResponse struct {
	Code      int           // the HTTP response code from WriteHeader
	Body      []byte
	HeaderMap http.Header   // the HTTP response headers
}

func respond(cached * CachedResponse, w http.ResponseWriter, wasCached bool) {
	for k, vs := range cached.HeaderMap {
		for _, v := range vs {
			// make-fmt.Println("key=", k, "value=", v)
			w.Header().Add(k, v)
		}
	}
	if wasCached {
		w.Header().Add("cache-status", "HIT")
	} else {
		w.Header().Add("cache-status", "MISS")
	}
	w.Write(cached.Body)
	w.WriteHeader(cached.Code)
}

func (h CacheHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) (int, error) {
	val, err := h.Client.Get(r.URL.Path).Result()
	if err != nil && err != redis.Nil {
		fmt.Println(err)
		return 500, err
	}

	if err == redis.Nil {
		rec := httptest.NewRecorder()
		status, err := h.Next.ServeHTTP(rec, r)

		cached := CachedResponse {
			Body: rec.Body.Bytes(),
			HeaderMap: rec.HeaderMap,
			Code: rec.Code,
		}

		serialized, err := msgpack.Marshal(cached)

		if err != nil {
			return 500, nil
		}

		err = h.Client.Set(r.URL.Path, serialized, time.Duration(5) * time.Minute).Err()
		if err != nil {
			return 500, err
		}

		respond(&cached, w, false)
		return status, err
	} else {
		cached := CachedResponse{}
		msgpack.Unmarshal([]byte(val), &cached)
		respond(&cached, w, true)
		return cached.Code, nil
	}
}
<commit_msg>Store page in cache only if headers allows to and respect max-age<commit_after>package cache

import (
	"fmt"
	"time"
	"net/http"
	"net/http/httptest"
	"gopkg.in/redis.v4"
	"gopkg.in/vmihailenco/msgpack.v2"
	"github.com/mholt/caddy/caddyhttp/httpserver"
	"github.com/pquerna/cachecontrol/cacheobject"
)

type CacheHandler struct {
	Client *redis.Client
	Next   httpserver.Handler
}

type CachedResponse struct {
	Code      int           // the HTTP response code from WriteHeader
	Body      []byte
	HeaderMap http.Header   // the HTTP response headers
}

func respond(response * CachedResponse, w http.ResponseWriter) {
	for k, vs := range response.HeaderMap {
		for _, v := range vs {
			// make-fmt.Println("key=", k, "value=", v)
			w.Header().Add(k, v)
		}
	}
	w.Write(response.Body)
	w.WriteHeader(response.Code)
}

func shouldUseCache(r *http.Request) bool {
	if r.Method != "GET" && r.Method != "HEAD" {
		// Only cache Get and head request
		return false
	}

	// Add more logic like get params, ?nocache=true
	return true
}

func getCacheableStatus(req *http.Request, res *httptest.ResponseRecorder) *cacheobject.ObjectResults {
	reqDir, _ := cacheobject.ParseRequestCacheControl(req.Header.Get("Cache-Control"))
	resDir, _ := cacheobject.ParseResponseCacheControl(res.Header().Get("Cache-Control"))
	expiresHeader, _ := http.ParseTime(res.Header().Get("Expires"))
	dateHeader, _ := http.ParseTime(res.Header().Get("Date"))
	lastModifiedHeader, _ := http.ParseTime(res.Header().Get("Last-Modified"))

	obj := cacheobject.Object{
		RespDirectives:         resDir,
		RespHeaders:            res.Header(),
		RespStatusCode:         res.Code,
		RespExpiresHeader:      expiresHeader,
		RespDateHeader:         dateHeader,
		RespLastModifiedHeader: lastModifiedHeader,

		ReqDirectives: reqDir,
		ReqHeaders:    req.Header,
		ReqMethod:     req.Method,

		NowUTC: time.Now().UTC(),
	}

	rv := cacheobject.ObjectResults{}
	cacheobject.CachableObject(&obj, &rv)
	cacheobject.ExpirationObject(&obj, &rv)
	return &rv
}

func isCacheable(rv *cacheobject.ObjectResults) bool {
	return len(rv.OutReasons) == 0 && getTTL(rv) > 0
}

func getTTL(rv *cacheobject.ObjectResults) time.Duration {
	return rv.OutExpirationTime.Sub(time.Now().UTC())
}

func getKey(r *http.Request) string {
	return "contentcache:" + r.Host + ":" + r.Method + ":" + r.URL.Path
}

func (h CacheHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) (int, error) {

	if !shouldUseCache(r) {
		return h.Next.ServeHTTP(w, r)
	}

	cached, err := h.Client.Get(getKey(r)).Result()
	if err != nil && err != redis.Nil {
		fmt.Println(err)
		return 500, err
	}

	if err == redis.Nil {
		rec := httptest.NewRecorder()
		status, err := h.Next.ServeHTTP(rec, r)

		response := CachedResponse {
			Body: rec.Body.Bytes(),
			HeaderMap: rec.HeaderMap,
			Code: rec.Code,
		}

		cacheableStatus := getCacheableStatus(r, rec)
		if isCacheable(cacheableStatus) {
			nextCache, err := msgpack.Marshal(response)

			if err != nil {
				return 500, nil
			}

			err = h.Client.Set(getKey(r), nextCache, getTTL(cacheableStatus)).Err()
			if err != nil {
				return 500, err
			}
		}

		respond(&response, w)
		return status, err
	} else {
		cache := CachedResponse{}
		msgpack.Unmarshal([]byte(cached), &cache)
		respond(&cache, w)
		return cache.Code, nil
	}
}
<|endoftext|>


<commit_before>// Simple caching library with expiration capabilities
package cache2go

import (
	"errors"
	"sync"
	"time"
)

// Structure of an entry in the cache
// data contains the user-set value in the cache
type CacheEntry struct {
	sync.Mutex
	key            interface{}
	data           interface{}
	lifeSpan       time.Duration
	createdOn      time.Time
	accessedOn     time.Time

	// Callback method triggered right before removing the item from the cache
	aboutToExpire func(interface{})
}

// Structure of a table with items in the cache
type CacheTable struct {
	sync.RWMutex
	name            string
	items           map[interface{}]*CacheEntry
	cleanupTimer    *time.Timer
	cleanupInterval time.Duration

	// Callback method triggered when trying to load a non-existing key
	loadData func(interface{}) *CacheEntry
}

var (
	cache = make(map[string]*CacheTable)
	mutex sync.RWMutex
)

// Mark entry to be kept for another expireDuration period
func (entry *CacheEntry) KeepAlive() {
	entry.Lock()
	defer entry.Unlock()
	entry.accessedOn = time.Now()
}

// Returns this entry's expiration duration
func (entry *CacheEntry) LifeSpan() time.Duration {
	entry.Lock()
	defer entry.Unlock()
	return entry.lifeSpan
}

// Returns when this entry was last accessed
func (entry *CacheEntry) AccessedOn() time.Time {
	entry.Lock()
	defer entry.Unlock()
	return entry.accessedOn
}

// Returns when this entry was added to the cache
func (entry *CacheEntry) CreatedOn() time.Time {
	entry.Lock()
	defer entry.Unlock()
	return entry.createdOn
}

// Returns the value of this cached item
func (entry *CacheEntry) Data() interface{} {
	entry.Lock()
	defer entry.Unlock()
	return entry.data
}

// Returns the existing cache table with given name or creates a new one
// if the table does not exist yet
func Cache(table string) *CacheTable {
	mutex.RLock()
	t, ok := cache[table]
	mutex.RUnlock()

	if !ok {
		t = &CacheTable{
			name:  table,
			items: make(map[interface{}]*CacheEntry),
		}
		mutex.Lock()
		cache[table] = t
		mutex.Unlock()
	}

	return t
}

// Returns how many items are currently stored in the cache
func (table *CacheTable) CacheCount() int {
	table.RLock()
	defer table.RUnlock()

	return len(table.items)
}

// Configures a data-loader callback, which will be called when trying
// to use access a non-existing key
func (table *CacheTable) SetDataLoader(f func(interface{}) *CacheEntry) {
	table.loadData = f
}

// Expiration check loop, triggered by a self-adjusting timer
func (table *CacheTable) expirationCheck() {
	table.Lock()
	if table.cleanupTimer != nil {
		table.cleanupTimer.Stop()
	}

	// Take a copy of cache so we can iterate over it without blocking the mutex
	cc := table.items
	table.Unlock()

	// To be more accurate with timers, we would need to update 'now' on every
	// loop iteration. Not sure it's really efficient though.
	now := time.Now()
	smallestDuration := 0 * time.Second
	for key, c := range cc {
		if c.lifeSpan == 0 {
			continue
		}
		if now.Sub(c.accessedOn) >= c.lifeSpan {
			if c.aboutToExpire != nil {
				c.aboutToExpire(key)
			}
			table.Lock()
			delete(table.items, key)
			table.Unlock()
		} else {
			if smallestDuration == 0 || c.lifeSpan < smallestDuration {
				smallestDuration = c.lifeSpan - now.Sub(c.accessedOn)
			}
		}
	}

	table.Lock()
	table.cleanupInterval = smallestDuration
	if smallestDuration > 0 {
		table.cleanupTimer = time.AfterFunc(smallestDuration, func() {
			go table.expirationCheck()
		})
	}
	table.Unlock()
}

/* Adds a key/value pair to the cache
 / key is a unique cache-item key in the cache
 / lifeSpan indicates how long this item will remain in the cache after its
 / last access
 / data is the cache-item value
 / The last parameter abouToExpireFunc can be nil. Otherwise abouToExpireFunc
 / will be called (with this item's key as its only parameter), right before
 / removing this item from the cache
*/
func (table *CacheTable) Cache(key interface{}, lifeSpan time.Duration, data interface{}, aboutToExpireFunc func(interface{})) {
	entry := CreateCacheEntry(key, lifeSpan, data, aboutToExpireFunc)

	table.Lock()
	table.items[key] = &entry
	expDur := table.cleanupInterval
	table.Unlock()

	// If we haven't set up any expiration check timer or found a more imminent item
	if lifeSpan > 0 && (expDur == 0 || lifeSpan < expDur) {
		table.expirationCheck()
	}
}

// Get an entry from the cache and mark it to be kept alive
func (table *CacheTable) Value(key interface{}) (*CacheEntry, error) {
	table.RLock()
	if r, ok := table.items[key]; ok {
		defer table.RUnlock()

		r.KeepAlive()
		return r, nil
	} else {
		table.RUnlock()

		if table.loadData != nil {
			item := table.loadData(key)
			table.Cache(key, item.lifeSpan, item.data, item.aboutToExpire)
			if item != nil {
				return item, nil
			} else {
				return nil, errors.New("Key not found and could not be loaded into cache")
			}
		}
	}

	return nil, errors.New("Key not found in cache")
}

// Returns a newly created CacheEntry
func CreateCacheEntry(key interface{}, lifeSpan time.Duration, data interface{}, aboutToExpireFunc func(interface{})) CacheEntry {
	t := time.Now()
	entry := CacheEntry{
		key: key,
		lifeSpan: lifeSpan,
		createdOn: t,
		accessedOn: t,
		aboutToExpire: aboutToExpireFunc,
		data: data,
	}

	return entry
}

// Delete all items from cache
func (table *CacheTable) Flush() {
	table.Lock()
	defer table.Unlock()

	table.items = make(map[interface{}]*CacheEntry)
	table.cleanupInterval = 0
	if table.cleanupTimer != nil {
		table.cleanupTimer.Stop()
	}

	mutex.Lock()
	defer mutex.Unlock()
	delete(cache, table.name)
}
<commit_msg>* Added callback method which gets triggered upon adding an item to the cache.<commit_after>// Simple caching library with expiration capabilities
package cache2go

import (
	"errors"
	"sync"
	"time"
)

// Structure of an entry in the cache
// data contains the user-set value in the cache
type CacheEntry struct {
	sync.Mutex
	key            interface{}
	data           interface{}
	lifeSpan       time.Duration
	createdOn      time.Time
	accessedOn     time.Time

	// Callback method triggered right before removing the item from the cache
	aboutToExpire func(interface{})
}

// Structure of a table with items in the cache
type CacheTable struct {
	sync.RWMutex
	name            string
	items           map[interface{}]*CacheEntry
	cleanupTimer    *time.Timer
	cleanupInterval time.Duration

	// Callback method triggered when trying to load a non-existing key
	loadData func(interface{}) *CacheEntry

	// Callback method triggered when adding a new item to the cache
	addedItem func(*CacheEntry)
}

var (
	cache = make(map[string]*CacheTable)
	mutex sync.RWMutex
)

// Mark entry to be kept for another expireDuration period
func (entry *CacheEntry) KeepAlive() {
	entry.Lock()
	defer entry.Unlock()
	entry.accessedOn = time.Now()
}

// Returns this entry's expiration duration
func (entry *CacheEntry) LifeSpan() time.Duration {
	entry.Lock()
	defer entry.Unlock()
	return entry.lifeSpan
}

// Returns when this entry was last accessed
func (entry *CacheEntry) AccessedOn() time.Time {
	entry.Lock()
	defer entry.Unlock()
	return entry.accessedOn
}

// Returns when this entry was added to the cache
func (entry *CacheEntry) CreatedOn() time.Time {
	entry.Lock()
	defer entry.Unlock()
	return entry.createdOn
}

// Returns the key of this cached item
func (entry *CacheEntry) Key() interface{} {
	entry.Lock()
	defer entry.Unlock()
	return entry.key
}

// Returns the value of this cached item
func (entry *CacheEntry) Data() interface{} {
	entry.Lock()
	defer entry.Unlock()
	return entry.data
}

// Returns the existing cache table with given name or creates a new one
// if the table does not exist yet
func Cache(table string) *CacheTable {
	mutex.RLock()
	t, ok := cache[table]
	mutex.RUnlock()

	if !ok {
		t = &CacheTable{
			name:  table,
			items: make(map[interface{}]*CacheEntry),
		}
		mutex.Lock()
		cache[table] = t
		mutex.Unlock()
	}

	return t
}

// Returns how many items are currently stored in the cache
func (table *CacheTable) CacheCount() int {
	table.RLock()
	defer table.RUnlock()

	return len(table.items)
}

// Configures a data-loader callback, which will be called when trying
// to use access a non-existing key
func (table *CacheTable) SetDataLoader(f func(interface{}) *CacheEntry) {
	table.loadData = f
}

// Configures a callback, which will be called every time a new item
// is added to the cache
func (table *CacheTable) SetAddedItemCallback(f func(*CacheEntry)) {
	table.addedItem = f
}

// Expiration check loop, triggered by a self-adjusting timer
func (table *CacheTable) expirationCheck() {
	table.Lock()
	if table.cleanupTimer != nil {
		table.cleanupTimer.Stop()
	}

	// Take a copy of cache so we can iterate over it without blocking the mutex
	cc := table.items
	table.Unlock()

	// To be more accurate with timers, we would need to update 'now' on every
	// loop iteration. Not sure it's really efficient though.
	now := time.Now()
	smallestDuration := 0 * time.Second
	for key, c := range cc {
		if c.lifeSpan == 0 {
			continue
		}
		if now.Sub(c.accessedOn) >= c.lifeSpan {
			if c.aboutToExpire != nil {
				c.aboutToExpire(key)
			}
			table.Lock()
			delete(table.items, key)
			table.Unlock()
		} else {
			if smallestDuration == 0 || c.lifeSpan < smallestDuration {
				smallestDuration = c.lifeSpan - now.Sub(c.accessedOn)
			}
		}
	}

	table.Lock()
	table.cleanupInterval = smallestDuration
	if smallestDuration > 0 {
		table.cleanupTimer = time.AfterFunc(smallestDuration, func() {
			go table.expirationCheck()
		})
	}
	table.Unlock()
}

/* Adds a key/value pair to the cache
 / key is a unique cache-item key in the cache
 / lifeSpan indicates how long this item will remain in the cache after its
 / last access
 / data is the cache-item value
 / The last parameter abouToExpireFunc can be nil. Otherwise abouToExpireFunc
 / will be called (with this item's key as its only parameter), right before
 / removing this item from the cache
*/
func (table *CacheTable) Cache(key interface{}, lifeSpan time.Duration, data interface{}, aboutToExpireFunc func(interface{})) {
	entry := CreateCacheEntry(key, lifeSpan, data, aboutToExpireFunc)

	table.Lock()
	table.items[key] = &entry
	expDur := table.cleanupInterval
	table.Unlock()

	// Trigger callback after adding an item to cache
	if table.addedItem != nil {
		table.addedItem(&entry)
	}

	// If we haven't set up any expiration check timer or found a more imminent item
	if lifeSpan > 0 && (expDur == 0 || lifeSpan < expDur) {
		table.expirationCheck()
	}
}

// Get an entry from the cache and mark it to be kept alive
func (table *CacheTable) Value(key interface{}) (*CacheEntry, error) {
	table.RLock()
	if r, ok := table.items[key]; ok {
		defer table.RUnlock()

		r.KeepAlive()
		return r, nil
	} else {
		table.RUnlock()

		if table.loadData != nil {
			item := table.loadData(key)
			table.Cache(key, item.lifeSpan, item.data, item.aboutToExpire)
			if item != nil {
				return item, nil
			} else {
				return nil, errors.New("Key not found and could not be loaded into cache")
			}
		}
	}

	return nil, errors.New("Key not found in cache")
}

// Returns a newly created CacheEntry
func CreateCacheEntry(key interface{}, lifeSpan time.Duration, data interface{}, aboutToExpireFunc func(interface{})) CacheEntry {
	t := time.Now()
	entry := CacheEntry{
		key: key,
		lifeSpan: lifeSpan,
		createdOn: t,
		accessedOn: t,
		aboutToExpire: aboutToExpireFunc,
		data: data,
	}

	return entry
}

// Delete all items from cache
func (table *CacheTable) Flush() {
	table.Lock()
	defer table.Unlock()

	table.items = make(map[interface{}]*CacheEntry)
	table.cleanupInterval = 0
	if table.cleanupTimer != nil {
		table.cleanupTimer.Stop()
	}

	mutex.Lock()
	defer mutex.Unlock()
	delete(cache, table.name)
}
<|endoftext|>