<commit_before>package configs

import (
	"bytes"
	"encoding/json"
	"os/exec"
)

type Rlimit struct {
	Type int    `json:"type"`
	Hard uint64 `json:"hard"`
	Soft uint64 `json:"soft"`
}

// IDMap represents UID/GID Mappings for User Namespaces.
type IDMap struct {
	ContainerID int `json:"container_id"`
	HostID      int `json:"host_id"`
	Size        int `json:"size"`
}

// Seccomp represents syscall restrictions
// By default, only the native architecture of the kernel is allowed to be used
// for syscalls. Additional architectures can be added by specifying them in
// Architectures.
type Seccomp struct {
	DefaultAction Action     `json:"default_action"`
	Architectures []string   `json:"architectures"`
	Syscalls      []*Syscall `json:"syscalls"`
}

// An action to be taken upon rule match in Seccomp
type Action int

const (
	Kill Action = iota + 1
	Errno
	Trap
	Allow
	Trace
)

// A comparison operator to be used when matching syscall arguments in Seccomp
type Operator int

const (
	EqualTo Operator = iota + 1
	NotEqualTo
	GreaterThan
	GreaterThanOrEqualTo
	LessThan
	LessThanOrEqualTo
	MaskEqualTo
)

// A rule to match a specific syscall argument in Seccomp
type Arg struct {
	Index    uint     `json:"index"`
	Value    uint64   `json:"value"`
	ValueTwo uint64   `json:"value_two"`
	Op       Operator `json:"op"`
}

// An rule to match a syscall in Seccomp
type Syscall struct {
	Name   string `json:"name"`
	Action Action `json:"action"`
	Args   []*Arg `json:"args"`
}

// TODO Windows. Many of these fields should be factored out into those parts
// which are common across platforms, and those which are platform specific.

// Config defines configuration options for executing a process inside a contained environment.
type Config struct {
	// NoPivotRoot will use MS_MOVE and a chroot to jail the process into the container's rootfs
	// This is a common option when the container is running in ramdisk
	NoPivotRoot bool `json:"no_pivot_root"`

	// ParentDeathSignal specifies the signal that is sent to the container's process in the case
	// that the parent process dies.
	ParentDeathSignal int `json:"parent_death_signal"`

	// PivotDir allows a custom directory inside the container's root filesystem to be used as pivot, when NoPivotRoot is not set.
	// When a custom PivotDir not set, a temporary dir inside the root filesystem will be used. The pivot dir needs to be writeable.
	// This is required when using read only root filesystems. In these cases, a read/writeable path can be (bind) mounted somewhere inside the root filesystem to act as pivot.
	PivotDir string `json:"pivot_dir"`

	// Path to a directory containing the container's root filesystem.
	Rootfs string `json:"rootfs"`

	// Readonlyfs will remount the container's rootfs as readonly where only externally mounted
	// bind mounts are writtable.
	Readonlyfs bool `json:"readonlyfs"`

	// Specifies the mount propagation flags to be applied to /.
	RootPropagation int `json:"rootPropagation"`

	// Mounts specify additional source and destination paths that will be mounted inside the container's
	// rootfs and mount namespace if specified
	Mounts []*Mount `json:"mounts"`

	// The device nodes that should be automatically created within the container upon container start.  Note, make sure that the node is marked as allowed in the cgroup as well!
	Devices []*Device `json:"devices"`

	MountLabel string `json:"mount_label"`

	// Hostname optionally sets the container's hostname if provided
	Hostname string `json:"hostname"`

	// Namespaces specifies the container's namespaces that it should setup when cloning the init process
	// If a namespace is not provided that namespace is shared from the container's parent process
	Namespaces Namespaces `json:"namespaces"`

	// Capabilities specify the capabilities to keep when executing the process inside the container
	// All capbilities not specified will be dropped from the processes capability mask
	Capabilities []string `json:"capabilities"`

	// Networks specifies the container's network setup to be created
	Networks []*Network `json:"networks"`

	// Routes can be specified to create entries in the route table as the container is started
	Routes []*Route `json:"routes"`

	// Cgroups specifies specific cgroup settings for the various subsystems that the container is
	// placed into to limit the resources the container has available
	Cgroups *Cgroup `json:"cgroups"`

	// AppArmorProfile specifies the profile to apply to the process running in the container and is
	// change at the time the process is execed
	AppArmorProfile string `json:"apparmor_profile"`

	// ProcessLabel specifies the label to apply to the process running in the container.  It is
	// commonly used by selinux
	ProcessLabel string `json:"process_label"`

	// Rlimits specifies the resource limits, such as max open files, to set in the container
	// If Rlimits are not set, the container will inherit rlimits from the parent process
	Rlimits []Rlimit `json:"rlimits"`

	// OomScoreAdj specifies the adjustment to be made by the kernel when calculating oom scores
	// for a process. Valid values are between the range [-1000, '1000'], where processes with
	// higher scores are preferred for being killed.
	// More information about kernel oom score calculation here: https://lwn.net/Articles/317814/
	OomScoreAdj int `json:"oom_score_adj"`

	// AdditionalGroups specifies the gids that should be added to supplementary groups
	// in addition to those that the user belongs to.
	AdditionalGroups []string `json:"additional_groups"`

	// UidMappings is an array of User ID mappings for User Namespaces
	UidMappings []IDMap `json:"uid_mappings"`

	// GidMappings is an array of Group ID mappings for User Namespaces
	GidMappings []IDMap `json:"gid_mappings"`

	// MaskPaths specifies paths within the container's rootfs to mask over with a bind
	// mount pointing to /dev/null as to prevent reads of the file.
	MaskPaths []string `json:"mask_paths"`

	// ReadonlyPaths specifies paths within the container's rootfs to remount as read-only
	// so that these files prevent any writes.
	ReadonlyPaths []string `json:"readonly_paths"`

	// Sysctl is a map of properties and their values. It is the equivalent of using
	// sysctl -w my.property.name value in Linux.
	Sysctl map[string]string `json:"sysctl"`

	// Seccomp allows actions to be taken whenever a syscall is made within the container.
	// A number of rules are given, each having an action to be taken if a syscall matches it.
	// A default action to be taken if no rules match is also given.
	Seccomp *Seccomp `json:"seccomp"`

	// NoNewPrivileges controls whether processes in the container can gain additional privileges.
	NoNewPrivileges bool `json:"no_new_privileges"`

	// Hooks are a collection of actions to perform at various container lifecycle events.
	// Hooks are not able to be marshaled to json but they are also not needed to.
	Hooks *Hooks `json:"-"`

	// Version is the version of opencontainer specification that is supported.
	Version string `json:"version"`
}

type Hooks struct {
	// Prestart commands are executed after the container namespaces are created,
	// but before the user supplied command is executed from init.
	Prestart []Hook

	// Poststart commands are executed after the container init process starts.
	Poststart []Hook

	// Poststop commands are executed after the container init process exits.
	Poststop []Hook
}

// HookState is the payload provided to a hook on execution.
type HookState struct {
	Version string `json:"version"`
	ID      string `json:"id"`
	Pid     int    `json:"pid"`
	Root    string `json:"root"`
}

type Hook interface {
	// Run executes the hook with the provided state.
	Run(HookState) error
}

// NewFunctionHooks will call the provided function when the hook is run.
func NewFunctionHook(f func(HookState) error) FuncHook {
	return FuncHook{
		run: f,
	}
}

type FuncHook struct {
	run func(HookState) error
}

func (f FuncHook) Run(s HookState) error {
	return f.run(s)
}

type Command struct {
	Path string   `json:"path"`
	Args []string `json:"args"`
	Env  []string `json:"env"`
	Dir  string   `json:"dir"`
}

// NewCommandHooks will execute the provided command when the hook is run.
func NewCommandHook(cmd Command) CommandHook {
	return CommandHook{
		Command: cmd,
	}
}

type CommandHook struct {
	Command
}

func (c Command) Run(s HookState) error {
	b, err := json.Marshal(s)
	if err != nil {
		return err
	}
	cmd := exec.Cmd{
		Path:  c.Path,
		Args:  c.Args,
		Env:   c.Env,
		Stdin: bytes.NewReader(b),
	}
	return cmd.Run()
}
<commit_msg>Add labels to libconatiner config<commit_after>package configs

import (
	"bytes"
	"encoding/json"
	"os/exec"
)

type Rlimit struct {
	Type int    `json:"type"`
	Hard uint64 `json:"hard"`
	Soft uint64 `json:"soft"`
}

// IDMap represents UID/GID Mappings for User Namespaces.
type IDMap struct {
	ContainerID int `json:"container_id"`
	HostID      int `json:"host_id"`
	Size        int `json:"size"`
}

// Seccomp represents syscall restrictions
// By default, only the native architecture of the kernel is allowed to be used
// for syscalls. Additional architectures can be added by specifying them in
// Architectures.
type Seccomp struct {
	DefaultAction Action     `json:"default_action"`
	Architectures []string   `json:"architectures"`
	Syscalls      []*Syscall `json:"syscalls"`
}

// An action to be taken upon rule match in Seccomp
type Action int

const (
	Kill Action = iota + 1
	Errno
	Trap
	Allow
	Trace
)

// A comparison operator to be used when matching syscall arguments in Seccomp
type Operator int

const (
	EqualTo Operator = iota + 1
	NotEqualTo
	GreaterThan
	GreaterThanOrEqualTo
	LessThan
	LessThanOrEqualTo
	MaskEqualTo
)

// A rule to match a specific syscall argument in Seccomp
type Arg struct {
	Index    uint     `json:"index"`
	Value    uint64   `json:"value"`
	ValueTwo uint64   `json:"value_two"`
	Op       Operator `json:"op"`
}

// An rule to match a syscall in Seccomp
type Syscall struct {
	Name   string `json:"name"`
	Action Action `json:"action"`
	Args   []*Arg `json:"args"`
}

// TODO Windows. Many of these fields should be factored out into those parts
// which are common across platforms, and those which are platform specific.

// Config defines configuration options for executing a process inside a contained environment.
type Config struct {
	// NoPivotRoot will use MS_MOVE and a chroot to jail the process into the container's rootfs
	// This is a common option when the container is running in ramdisk
	NoPivotRoot bool `json:"no_pivot_root"`

	// ParentDeathSignal specifies the signal that is sent to the container's process in the case
	// that the parent process dies.
	ParentDeathSignal int `json:"parent_death_signal"`

	// PivotDir allows a custom directory inside the container's root filesystem to be used as pivot, when NoPivotRoot is not set.
	// When a custom PivotDir not set, a temporary dir inside the root filesystem will be used. The pivot dir needs to be writeable.
	// This is required when using read only root filesystems. In these cases, a read/writeable path can be (bind) mounted somewhere inside the root filesystem to act as pivot.
	PivotDir string `json:"pivot_dir"`

	// Path to a directory containing the container's root filesystem.
	Rootfs string `json:"rootfs"`

	// Readonlyfs will remount the container's rootfs as readonly where only externally mounted
	// bind mounts are writtable.
	Readonlyfs bool `json:"readonlyfs"`

	// Specifies the mount propagation flags to be applied to /.
	RootPropagation int `json:"rootPropagation"`

	// Mounts specify additional source and destination paths that will be mounted inside the container's
	// rootfs and mount namespace if specified
	Mounts []*Mount `json:"mounts"`

	// The device nodes that should be automatically created within the container upon container start.  Note, make sure that the node is marked as allowed in the cgroup as well!
	Devices []*Device `json:"devices"`

	MountLabel string `json:"mount_label"`

	// Hostname optionally sets the container's hostname if provided
	Hostname string `json:"hostname"`

	// Namespaces specifies the container's namespaces that it should setup when cloning the init process
	// If a namespace is not provided that namespace is shared from the container's parent process
	Namespaces Namespaces `json:"namespaces"`

	// Capabilities specify the capabilities to keep when executing the process inside the container
	// All capbilities not specified will be dropped from the processes capability mask
	Capabilities []string `json:"capabilities"`

	// Networks specifies the container's network setup to be created
	Networks []*Network `json:"networks"`

	// Routes can be specified to create entries in the route table as the container is started
	Routes []*Route `json:"routes"`

	// Cgroups specifies specific cgroup settings for the various subsystems that the container is
	// placed into to limit the resources the container has available
	Cgroups *Cgroup `json:"cgroups"`

	// AppArmorProfile specifies the profile to apply to the process running in the container and is
	// change at the time the process is execed
	AppArmorProfile string `json:"apparmor_profile"`

	// ProcessLabel specifies the label to apply to the process running in the container.  It is
	// commonly used by selinux
	ProcessLabel string `json:"process_label"`

	// Rlimits specifies the resource limits, such as max open files, to set in the container
	// If Rlimits are not set, the container will inherit rlimits from the parent process
	Rlimits []Rlimit `json:"rlimits"`

	// OomScoreAdj specifies the adjustment to be made by the kernel when calculating oom scores
	// for a process. Valid values are between the range [-1000, '1000'], where processes with
	// higher scores are preferred for being killed.
	// More information about kernel oom score calculation here: https://lwn.net/Articles/317814/
	OomScoreAdj int `json:"oom_score_adj"`

	// AdditionalGroups specifies the gids that should be added to supplementary groups
	// in addition to those that the user belongs to.
	AdditionalGroups []string `json:"additional_groups"`

	// UidMappings is an array of User ID mappings for User Namespaces
	UidMappings []IDMap `json:"uid_mappings"`

	// GidMappings is an array of Group ID mappings for User Namespaces
	GidMappings []IDMap `json:"gid_mappings"`

	// MaskPaths specifies paths within the container's rootfs to mask over with a bind
	// mount pointing to /dev/null as to prevent reads of the file.
	MaskPaths []string `json:"mask_paths"`

	// ReadonlyPaths specifies paths within the container's rootfs to remount as read-only
	// so that these files prevent any writes.
	ReadonlyPaths []string `json:"readonly_paths"`

	// Sysctl is a map of properties and their values. It is the equivalent of using
	// sysctl -w my.property.name value in Linux.
	Sysctl map[string]string `json:"sysctl"`

	// Seccomp allows actions to be taken whenever a syscall is made within the container.
	// A number of rules are given, each having an action to be taken if a syscall matches it.
	// A default action to be taken if no rules match is also given.
	Seccomp *Seccomp `json:"seccomp"`

	// NoNewPrivileges controls whether processes in the container can gain additional privileges.
	NoNewPrivileges bool `json:"no_new_privileges"`

	// Hooks are a collection of actions to perform at various container lifecycle events.
	// Hooks are not able to be marshaled to json but they are also not needed to.
	Hooks *Hooks `json:"-"`

	// Version is the version of opencontainer specification that is supported.
	Version string `json:"version"`

	// Labels are user defined metadata that is stored in the config and populated on the state
	Labels []string `json:"labels"`
}

type Hooks struct {
	// Prestart commands are executed after the container namespaces are created,
	// but before the user supplied command is executed from init.
	Prestart []Hook

	// Poststart commands are executed after the container init process starts.
	Poststart []Hook

	// Poststop commands are executed after the container init process exits.
	Poststop []Hook
}

// HookState is the payload provided to a hook on execution.
type HookState struct {
	Version string `json:"version"`
	ID      string `json:"id"`
	Pid     int    `json:"pid"`
	Root    string `json:"root"`
}

type Hook interface {
	// Run executes the hook with the provided state.
	Run(HookState) error
}

// NewFunctionHooks will call the provided function when the hook is run.
func NewFunctionHook(f func(HookState) error) FuncHook {
	return FuncHook{
		run: f,
	}
}

type FuncHook struct {
	run func(HookState) error
}

func (f FuncHook) Run(s HookState) error {
	return f.run(s)
}

type Command struct {
	Path string   `json:"path"`
	Args []string `json:"args"`
	Env  []string `json:"env"`
	Dir  string   `json:"dir"`
}

// NewCommandHooks will execute the provided command when the hook is run.
func NewCommandHook(cmd Command) CommandHook {
	return CommandHook{
		Command: cmd,
	}
}

type CommandHook struct {
	Command
}

func (c Command) Run(s HookState) error {
	b, err := json.Marshal(s)
	if err != nil {
		return err
	}
	cmd := exec.Cmd{
		Path:  c.Path,
		Args:  c.Args,
		Env:   c.Env,
		Stdin: bytes.NewReader(b),
	}
	return cmd.Run()
}
<|endoftext|>


<commit_before>package dom

/*
 * Implements a very small, very non-compliant subset of the DOM Core Level 3
 * http://www.w3.org/TR/DOM-Level-3-Core/
 *
 * Copyright (c) 2009, Rob Russell
 * Copyright (c) 2009, Jeff Schiller
 */

import (
  "container/vector";
  "strings";
  "xml";
  "fmt";
  "os";
)

type _node struct {
  p *_node; // parent
  c vector.Vector; // children
}

func (n *_node) NodeName() string { return "Node.NodeName() not implemented"; }
func (n *_node) NodeType() int { return -1; }
func (n *_node) AppendChild(child Node) Node {
  n.c.Push(child);
  return child;
}

func newNode() (n *_node) {
  n = new(_node);
  return;
}

// implements the Element interface
type _elem struct {
  *_node;
  n xml.Name; // name
}
func (e *_elem) NodeName() string { return e.n.Local; }
func (e *_elem) NodeType() int { return 1; }
func (e *_elem) TagName() string { return e.NodeName(); }
func (e *_elem) GetAttribute(name string) string { return "Element.getAttribute() not implemented"; }
func (e *_elem) SetAttribute(attrname string, attrval string) { }

func newElem(token xml.StartElement) (e *_elem) {
  e = new(_elem);
  e.n = token.Name;
  return;
}

// implements the Document interface
type _doc struct {
  *_node;
  root Element;
}
func (d *_doc) NodeName() string { return "#document"; }
func (d *_doc) NodeType() int { return 9; }
func (d *_doc) DocumentElement() Element { return d.root; }
func (d *_doc) setRoot(r Element) Element {
  d.root = r;
  return r;
}

func ParseString(s string) Document {
  r := strings.NewReader(s);
  p := xml.NewParser(r);
  t, err := p.Token();
  d := new(_doc);
  e := (Element)(nil); // e is the current parent
  for t != nil {
    switch token := t.(type) {
      case xml.StartElement:
        newElem := newElem(token);
        if e == nil {
          // set doc root
          e = d.setRoot(newElem);
        } else {
          // this element is a child of e, the last element we found
          e,_ = e.AppendChild(newElem).(Element);
        }
      case xml.EndElement:
      	// TODO: go up to parent
      default:
        fmt.Println("Unknown type");
    }
    // get the next token
    t, err = p.Token();
  }
  if err != os.EOF {
    fmt.Println(err.String());
  }
  return d;
}
<commit_msg>Added Element getAttribute and setAttribute methods.<commit_after>package dom

/*
 * Implements a very small, very non-compliant subset of the DOM Core Level 3
 * http://www.w3.org/TR/DOM-Level-3-Core/
 *
 * Copyright (c) 2009, Rob Russell
 * Copyright (c) 2009, Jeff Schiller
 */

import (
  "container/vector";
  "strings";
  "xml";
  "fmt";
  "os";
)

type _node struct {
  p *_node; // parent
  c vector.Vector; // children
}

func (n *_node) NodeName() string { return "Node.NodeName() not implemented"; }
func (n *_node) NodeType() int { return -1; }
func (n *_node) AppendChild(child Node) Node {
  n.c.Push(child);
  return child;
}

func newNode() (n *_node) {
  n = new(_node);
  return;
}

// implements the Element interface
type _elem struct {
  *_node;
  n xml.Name; // name
  attribs map[string] string; // attributes of the element
}
func (e *_elem) NodeName() string { return e.n.Local; }
func (e *_elem) NodeType() int { return 1; }
func (e *_elem) TagName() string { return e.NodeName(); }
func (e *_elem) GetAttribute(name string) string {
  val, ok := e.attribs[name];
  if (!ok) {
    val = "";
  }
  return val;
}
func (e *_elem) SetAttribute(attrname string, attrval string) {
  e.attribs[attrname]=attrval;
}

func newElem(token xml.StartElement) (e *_elem) {
  e = new(_elem);
  e.n = token.Name;
  e.attribs = make(map[string] string);
  return;
}

// implements the Document interface
type _doc struct {
  *_node;
  root Element;
}
func (d *_doc) NodeName() string { return "#document"; }
func (d *_doc) NodeType() int { return 9; }
func (d *_doc) DocumentElement() Element { return d.root; }
func (d *_doc) setRoot(r Element) Element {
  d.root = r;
  return r;
}

func ParseString(s string) Document {
  r := strings.NewReader(s);
  p := xml.NewParser(r);
  t, err := p.Token();
  d := new(_doc);
  e := (Element)(nil); // e is the current parent
  for t != nil {
    switch token := t.(type) {
      case xml.StartElement:
        el := newElem(token);
        for ar := range(token.Attr) {
          //fmt.Println("ar ",ar," attr ",token.Attr[ar].Name.Local);
          el.SetAttribute(token.Attr[ar].Name.Local, token.Attr[ar].Value);
        }
        if e == nil {
          // set doc root
          e = d.setRoot(el);
        } else {
          // this element is a child of e, the last element we found
          e,_ = e.AppendChild(el).(Element);
        }
      case xml.EndElement:
      	// TODO: go up to parent
      default:
        fmt.Println("Unknown type");
    }
    // get the next token
    t, err = p.Token();
  }
  if err != os.EOF {
    fmt.Println(err.String());
  }
  return d;
}
<|endoftext|>


<commit_before>// Copyright 2013 Am Laher.
// This code is adapted from code within the Go tree.
// See Go's licence information below:
//
// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package ar

import (
	"errors"
	"fmt"
	"io"
	"strings"
)

var (
	//ErrWriteAfterClose shows that a write was attempted after the archive has been closed (and the footer is written)
	ErrWriteAfterClose = errors.New("ar: write after close")
	errNameTooLong     = errors.New("ar: name too long")
	errInvalidHeader   = errors.New("ar: header field too long or contains invalid values")
)

// A Writer provides sequential writing of an ar archive.
// An ar archive consists of a sequence of files.
// Call WriteHeader to begin a new file, and then call Write to supply that file's data,
// writing at most hdr.Size bytes in total.
type Writer struct {
	w                       io.Writer
	arFileHeaderWritten     bool
	err                     error
	nb                      int64 // number of unwritten bytes for current file entry
	pad                     bool  // whether the file will be padded an extra byte (i.e. if ther's an odd number of bytes in the file)
	closed                  bool
	TerminateFilenamesSlash bool
}

// NewWriter creates a new Writer writing to w.
func NewWriter(w io.Writer) *Writer { return &Writer{w: w} }

// Flush finishes writing the current file (optional. This is called by writeHeader anyway.)
func (aw *Writer) Flush() error {
	if aw.nb > 0 {
		aw.err = fmt.Errorf("ar: missed writing %d bytes", aw.nb)
		return aw.err
	}
	if !aw.arFileHeaderWritten {
		_, aw.err = aw.w.Write([]byte(ArFileHeader))
		if aw.err != nil {
			return aw.err
		}
		aw.arFileHeaderWritten = true
	}
	if aw.pad {
		//pad with a newline
		if _, aw.err = aw.w.Write([]byte("\n")); aw.err != nil {
			return aw.err
		}
	}
	aw.nb = 0
	aw.pad = false
	return aw.err
}

// WriteHeader writes hdr and prepares to accept the file's contents.
// WriteHeader calls Flush if it is not the first header.
// Calling after a Close will return ErrWriteAfterClose.
func (aw *Writer) WriteHeader(hdr *Header) error {
	return aw.writeHeader(hdr)
}

// WriteHeader writes hdr and prepares to accept the file's contents.
// WriteHeader calls Flush if it is not the first header.
// Calling after a Close will return ErrWriteAfterClose.
func (aw *Writer) writeHeader(hdr *Header) error {
	if aw.closed {
		return ErrWriteAfterClose
	}
	if aw.err == nil {
		aw.Flush()
	}
	if aw.err != nil {
		return aw.err
	}
	fmodTimestamp := fmt.Sprintf("%d", hdr.ModTime.Unix())
	//use root by default (this is particularly useful for debs).
	uid := fmt.Sprintf("%d", hdr.Uid)
	if len(uid) > 6 {
		return fmt.Errorf("UID too long")
	}
	gid := fmt.Sprintf("%d", hdr.Gid)
	if len(gid) > 6 {
		return fmt.Errorf("GID too long")
	}
	//Files only atm (not dirs)
	mode := fmt.Sprintf("100%d", hdr.Mode)
	size := fmt.Sprintf("%d", hdr.Size)
	name := hdr.Name
	if aw.TerminateFilenamesSlash {
		name += "/"
	}
	line := fmt.Sprintf("%s%s%s%s%s%s`\n", pad(name, 16), pad(fmodTimestamp, 12), pad(gid, 6), pad(uid, 6), pad(mode, 8), pad(size, 10))
	if _, err := aw.Write([]byte(line)); err != nil {
		return err
	}
	// data section is 2-byte aligned.
	if hdr.Size%2 == 1 {
		aw.pad = true
	}
	aw.nb = hdr.Size
	return nil
}

// Write some data to the ar file.
func (aw *Writer) Write(b []byte) (int, error) {
	if aw.closed {
		aw.err = ErrWriteAfterClose
		return 0, aw.err
	}
	var n int
	n, aw.err = aw.w.Write(b)
	if aw.err != nil {
		return n, aw.err
	}
	aw.nb -= int64(n)
	return n, aw.err
}

// Close closes the ar archive, flushing any unwritten
// data to the underlying writer.
func (aw *Writer) Close() error {
	if aw.err != nil || aw.closed {
		return aw.err
	}
	aw.Flush()
	aw.closed = true
	if aw.err != nil {
		return aw.err
	}
	return aw.err
}

// pads a value with spaces up to a given length
func pad(value string, length int) string {
	plen := length - len(value)
	if plen > 0 {
		return value + strings.Repeat(" ", plen)
	}
	return value
}
<commit_msg>documentation for TerminateFilenamesSlash<commit_after>// Copyright 2013 Am Laher.
// This code is adapted from code within the Go tree.
// See Go's licence information below:
//
// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package ar

import (
	"errors"
	"fmt"
	"io"
	"strings"
)

var (
	//ErrWriteAfterClose shows that a write was attempted after the archive has been closed (and the footer is written)
	ErrWriteAfterClose = errors.New("ar: write after close")
	errNameTooLong     = errors.New("ar: name too long")
	errInvalidHeader   = errors.New("ar: header field too long or contains invalid values")
)

// A Writer provides sequential writing of an ar archive.
// An ar archive consists of a sequence of files.
// Call WriteHeader to begin a new file, and then call Write to supply that file's data,
// writing at most hdr.Size bytes in total.
type Writer struct {
	w                       io.Writer
	arFileHeaderWritten     bool
	err                     error
	nb                      int64 // number of unwritten bytes for current file entry
	pad                     bool  // whether the file will be padded an extra byte (i.e. if ther's an odd number of bytes in the file)
	closed                  bool
	TerminateFilenamesSlash bool // This flag determines whether to terminate filenames with a slash '/' or not. GNU ar uses slashes, whereas .deb files tend not to use them.
}

// NewWriter creates a new Writer writing to w.
func NewWriter(w io.Writer) *Writer { return &Writer{w: w} }

// Flush finishes writing the current file (optional. This is called by writeHeader anyway.)
func (aw *Writer) Flush() error {
	if aw.nb > 0 {
		aw.err = fmt.Errorf("ar: missed writing %d bytes", aw.nb)
		return aw.err
	}
	if !aw.arFileHeaderWritten {
		_, aw.err = aw.w.Write([]byte(ArFileHeader))
		if aw.err != nil {
			return aw.err
		}
		aw.arFileHeaderWritten = true
	}
	if aw.pad {
		//pad with a newline
		if _, aw.err = aw.w.Write([]byte("\n")); aw.err != nil {
			return aw.err
		}
	}
	aw.nb = 0
	aw.pad = false
	return aw.err
}

// WriteHeader writes hdr and prepares to accept the file's contents.
// WriteHeader calls Flush if it is not the first header.
// Calling after a Close will return ErrWriteAfterClose.
func (aw *Writer) WriteHeader(hdr *Header) error {
	return aw.writeHeader(hdr)
}

// WriteHeader writes hdr and prepares to accept the file's contents.
// WriteHeader calls Flush if it is not the first header.
// Calling after a Close will return ErrWriteAfterClose.
func (aw *Writer) writeHeader(hdr *Header) error {
	if aw.closed {
		return ErrWriteAfterClose
	}
	if aw.err == nil {
		aw.Flush()
	}
	if aw.err != nil {
		return aw.err
	}
	fmodTimestamp := fmt.Sprintf("%d", hdr.ModTime.Unix())
	//use root by default (this is particularly useful for debs).
	uid := fmt.Sprintf("%d", hdr.Uid)
	if len(uid) > 6 {
		return fmt.Errorf("UID too long")
	}
	gid := fmt.Sprintf("%d", hdr.Gid)
	if len(gid) > 6 {
		return fmt.Errorf("GID too long")
	}
	//Files only atm (not dirs)
	mode := fmt.Sprintf("100%d", hdr.Mode)
	size := fmt.Sprintf("%d", hdr.Size)
	name := hdr.Name
	if aw.TerminateFilenamesSlash {
		name += "/"
	}
	line := fmt.Sprintf("%s%s%s%s%s%s`\n", pad(name, 16), pad(fmodTimestamp, 12), pad(gid, 6), pad(uid, 6), pad(mode, 8), pad(size, 10))
	if _, err := aw.Write([]byte(line)); err != nil {
		return err
	}
	// data section is 2-byte aligned.
	if hdr.Size%2 == 1 {
		aw.pad = true
	}
	aw.nb = hdr.Size
	return nil
}

// Write some data to the ar file.
func (aw *Writer) Write(b []byte) (int, error) {
	if aw.closed {
		aw.err = ErrWriteAfterClose
		return 0, aw.err
	}
	var n int
	n, aw.err = aw.w.Write(b)
	if aw.err != nil {
		return n, aw.err
	}
	aw.nb -= int64(n)
	return n, aw.err
}

// Close closes the ar archive, flushing any unwritten
// data to the underlying writer.
func (aw *Writer) Close() error {
	if aw.err != nil || aw.closed {
		return aw.err
	}
	aw.Flush()
	aw.closed = true
	if aw.err != nil {
		return aw.err
	}
	return aw.err
}

// pads a value with spaces up to a given length
func pad(value string, length int) string {
	plen := length - len(value)
	if plen > 0 {
		return value + strings.Repeat(" ", plen)
	}
	return value
}
<|endoftext|>


<commit_before>package ast

import (
	"fmt"
	"strings"
)

// Index represents an indexing operation into another data structure
type Index struct {
	Target Node
	Key    Node
	Posx   Pos
}

func (n *Index) Accept(v Visitor) Node {
	return v(n)
}

func (n *Index) Pos() Pos {
	return n.Posx
}

func (n *Index) String() string {
	return fmt.Sprintf("Index(%s, %s)", n.Target, n.Key)
}

func (n *Index) Type(s Scope) (Type, error) {
	variableAccess, ok := n.Target.(*VariableAccess)
	if !ok {
		return TypeInvalid, fmt.Errorf("target is not a variable")
	}

	variable, ok := s.LookupVar(variableAccess.Name)
	if !ok {
		return TypeInvalid, fmt.Errorf("unknown variable accessed: %s", variableAccess.Name)
	}

	switch variable.Type {
	case TypeList:
		return n.typeList(variable, variableAccess.Name)
	case TypeMap:
		return n.typeMap(variable, variableAccess.Name)
	default:
		return TypeInvalid, fmt.Errorf("invalid index operation into non-indexable type: %s", variable.Type)
	}
}

func (n *Index) typeList(variable Variable, variableName string) (Type, error) {
	// We assume type checking has already determined that this is a list
	list := variable.Value.([]Variable)

	return VariableListElementTypesAreHomogenous(variableName, list)
}

func (n *Index) typeMap(variable Variable, variableName string) (Type, error) {
	// We assume type checking has already determined that this is a map
	vmap := variable.Value.(map[string]Variable)

	return VariableMapValueTypesAreHomogenous(variableName, vmap)
}

func reportTypes(typesFound map[Type]struct{}) string {
	stringTypes := make([]string, len(typesFound))
	i := 0
	for k, _ := range typesFound {
		stringTypes[0] = k.String()
		i++
	}
	return strings.Join(stringTypes, ", ")
}

func (n *Index) GoString() string {
	return fmt.Sprintf("*%#v", *n)
}
<commit_msg>Indexes may require further interpolation<commit_after>package ast

import (
	"fmt"
	"strings"
)

// Index represents an indexing operation into another data structure
type Index struct {
	Target Node
	Key    Node
	Posx   Pos
}

func (n *Index) Accept(v Visitor) Node {
	// the Key may have further interpolations
	switch n.Key.(type) {
	case *Call, *VariableAccess:
		n.Key = n.Key.Accept(v)
	}
	return v(n)
}

func (n *Index) Pos() Pos {
	return n.Posx
}

func (n *Index) String() string {
	return fmt.Sprintf("Index(%s, %s)", n.Target, n.Key)
}

func (n *Index) Type(s Scope) (Type, error) {
	variableAccess, ok := n.Target.(*VariableAccess)
	if !ok {
		return TypeInvalid, fmt.Errorf("target is not a variable")
	}

	variable, ok := s.LookupVar(variableAccess.Name)
	if !ok {
		return TypeInvalid, fmt.Errorf("unknown variable accessed: %s", variableAccess.Name)
	}

	switch variable.Type {
	case TypeList:
		return n.typeList(variable, variableAccess.Name)
	case TypeMap:
		return n.typeMap(variable, variableAccess.Name)
	default:
		return TypeInvalid, fmt.Errorf("invalid index operation into non-indexable type: %s", variable.Type)
	}
}

func (n *Index) typeList(variable Variable, variableName string) (Type, error) {
	// We assume type checking has already determined that this is a list
	list := variable.Value.([]Variable)

	return VariableListElementTypesAreHomogenous(variableName, list)
}

func (n *Index) typeMap(variable Variable, variableName string) (Type, error) {
	// We assume type checking has already determined that this is a map
	vmap := variable.Value.(map[string]Variable)

	return VariableMapValueTypesAreHomogenous(variableName, vmap)
}

func reportTypes(typesFound map[Type]struct{}) string {
	stringTypes := make([]string, len(typesFound))
	i := 0
	for k, _ := range typesFound {
		stringTypes[0] = k.String()
		i++
	}
	return strings.Join(stringTypes, ", ")
}

func (n *Index) GoString() string {
	return fmt.Sprintf("*%#v", *n)
}
<|endoftext|>


<commit_before>package mongo

import (
	// Standard Library Imports
	"context"
	"crypto/tls"
	"fmt"
	"net"
	"time"

	// External Imports
	"github.com/globalsign/mgo"
	"github.com/ory/fosite"
	"github.com/sirupsen/logrus"

	// Local Imports
	"github.com/matthewhartstonge/storage"
)

func init() {
	// Bind a logger, but only to panic level. Leave it to the user to decide
	// whether they want datastore logging or not.
	SetLogger(logrus.New())
	logger.Level = logrus.PanicLevel
}

const (
	defaultHost         = "localhost"
	defaultPort         = 27017
	defaultDatabaseName = "oauth2"
)

type Store struct {
	// Internals
	DB *mgo.Database

	// Public API
	Hasher fosite.Hasher
	storage.Store
}

// NewSession returns a mongo session.
// Note: The session requires closing manually so no memory leaks occur.
// This is best achieved by calling `defer session.Close()` straight after
// obtaining the returned session object.
func (m *Store) NewSession() (session *mgo.Session) {
	return m.DB.Session.Copy()
}

// Close terminates the mongo session.
func (m *Store) Close() {
	m.DB.Session.Close()
}

// Config defines the configuration parameters which are used by GetMongoSession.
type Config struct {
	Hostnames    []string `default:"localhost" envconfig:"CONNECTIONS_MONGO_HOSTNAMES"`
	Port         uint16   `default:"27017"     envconfig:"CONNECTIONS_MONGO_PORT"`
	AuthDB       string   `default:"admin"     envconfig:"CONNECTIONS_MONGO_AUTHDB"`
	Username     string   `default:""          envconfig:"CONNECTIONS_MONGO_USERNAME"`
	Password     string   `default:""          envconfig:"CONNECTIONS_MONGO_PASSWORD"`
	DatabaseName string   `default:""          envconfig:"CONNECTIONS_MONGO_NAME"`
	Replset      string   `default:""          envconfig:"CONNECTIONS_MONGO_REPLSET"`
	Timeout      uint     `default:"10"        envconfig:"CONNECTIONS_MONGO_TIMEOUT"`
	SSL          bool     `default:"false"     envconfig:"CONNECTIONS_MONGO_SSL"`
	TLSConfig    *tls.Config
}

// DefaultConfig returns a configuration for a locally hosted, unauthenticated mongo
func DefaultConfig() *Config {
	return &Config{
		Hostnames:    []string{defaultHost},
		Port:         defaultPort,
		DatabaseName: defaultDatabaseName,
	}
}

// ConnectionInfo configures options for establishing a session with a MongoDB cluster.
func ConnectionInfo(cfg *Config) *mgo.DialInfo {
	if len(cfg.Hostnames) == 0 {
		cfg.Hostnames = []string{defaultHost}
	}

	if cfg.DatabaseName == "" {
		cfg.DatabaseName = defaultDatabaseName
	}

	if cfg.Port > 0 {
		for i := range cfg.Hostnames {
			cfg.Hostnames[i] = fmt.Sprintf("%s:%d", cfg.Hostnames[i], cfg.Port)
		}
	}

	if cfg.Timeout == 0 {
		cfg.Timeout = 10
	}

	dialInfo := &mgo.DialInfo{
		Addrs:          cfg.Hostnames,
		Database:       cfg.DatabaseName,
		Username:       cfg.Username,
		Password:       cfg.Password,
		Source:         cfg.AuthDB,
		ReplicaSetName: cfg.Replset,
		Timeout:        time.Second * time.Duration(cfg.Timeout),
	}

	if cfg.SSL {
		dialInfo.DialServer = func(addr *mgo.ServerAddr) (net.Conn, error) {
			return tls.Dial("tcp", addr.String(), cfg.TLSConfig)
		}
	}

	return dialInfo
}

// Connect returns a connection to mongo.
func Connect(cfg *Config) (*mgo.Database, error) {
	log := logger.WithFields(logrus.Fields{
		"package": "mongo",
		"method":  "Connect",
	})

	dialInfo := ConnectionInfo(cfg)
	session, err := mgo.DialWithInfo(dialInfo)
	if err != nil {
		log.WithError(err).Error("Unable to connect to mongo! Have you configured your connection properly?")
		return nil, err
	}

	// Monotonic consistency will start reading from a slave if possible
	session.SetMode(mgo.Monotonic, true)
	return session.DB(cfg.DatabaseName), nil
}

// New allows for custom mongo configuration and custom hashers.
func New(cfg *Config, hashee fosite.Hasher) (*Store, error) {
	log := logger.WithFields(logrus.Fields{
		"package": "mongo",
		"method":  "NewFromConfig",
	})

	database, err := Connect(cfg)
	if err != nil {
		log.WithError(err).Error("Unable to connect to mongo! Are you sure mongo is running on localhost?")
		return nil, err
	}

	if hashee == nil {
		// Initialize default fosite Hasher.
		hashee = &fosite.BCrypt{
			WorkFactor: 10,
		}
	}

	// Build up the mongo endpoints
	mongoCache := &CacheManager{
		DB: database,
	}
	mongoClients := &ClientManager{
		DB:     database,
		Hasher: hashee,
	}
	mongoUsers := &UserManager{
		DB:     database,
		Hasher: hashee,
	}
	mongoRequests := &RequestManager{
		DB: database,

		Cache:   mongoCache,
		Clients: mongoClients,
		Users:   mongoUsers,
	}

	// Init DB collections, indices e.t.c.
	managers := []storage.Configurer{
		mongoCache,
		mongoClients,
		mongoUsers,
		mongoRequests,
	}

	// Use the main DB database to configure the mongo collections on first up.
	mgoSession := database.Session.Copy()
	defer mgoSession.Close()
	ctx := MgoSessionToContext(context.Background(), mgoSession)

	for _, manager := range managers {
		err := manager.Configure(ctx)
		if err != nil {
			log.WithError(err).Error("Unable to configure mongo collections!")
			return nil, err
		}
	}

	return &Store{
		DB:     database,
		Hasher: hashee,
		Store: storage.Store{
			CacheManager:   mongoCache,
			ClientManager:  mongoClients,
			RequestManager: mongoRequests,
			UserManager:    mongoUsers,
		},
	}, nil
}

// NewDefaultStore returns a Store configured with the default mongo
// configuration and default Hasher.
func NewDefaultStore() (*Store, error) {
	cfg := DefaultConfig()
	return New(cfg, nil)
}
<commit_msg>:zap: Fixes #20.<commit_after>package mongo

import (
	// Standard Library Imports
	"context"
	"crypto/tls"
	"fmt"
	"net"
	"time"

	// External Imports
	"github.com/globalsign/mgo"
	"github.com/ory/fosite"
	"github.com/sirupsen/logrus"

	// Local Imports
	"github.com/matthewhartstonge/storage"
)

func init() {
	// Bind a logger, but only to panic level. Leave it to the user to decide
	// whether they want datastore logging or not.
	SetLogger(logrus.New())
	logger.Level = logrus.PanicLevel
}

const (
	defaultHost         = "localhost"
	defaultPort         = 27017
	defaultDatabaseName = "oauth2"
)

type Store struct {
	// Internals
	DB *mgo.Database

	// Public API
	Hasher fosite.Hasher
	storage.Store
}

// NewSession returns a mongo session.
// Note: The session requires closing manually so no memory leaks occur.
// This is best achieved by calling `defer session.Close()` straight after
// obtaining the returned session object.
func (m *Store) NewSession() (session *mgo.Session) {
	return m.DB.Session.Copy()
}

// Close terminates the mongo session.
func (m *Store) Close() {
	m.DB.Session.Close()
}

// Config defines the configuration parameters which are used by GetMongoSession.
type Config struct {
	Hostnames    []string    `default:"localhost" envconfig:"CONNECTIONS_MONGO_HOSTNAMES"`
	Port         uint16      `default:"27017"     envconfig:"CONNECTIONS_MONGO_PORT"`
	AuthDB       string      `default:"admin"     envconfig:"CONNECTIONS_MONGO_AUTHDB"`
	Username     string      `default:""          envconfig:"CONNECTIONS_MONGO_USERNAME"`
	Password     string      `default:""          envconfig:"CONNECTIONS_MONGO_PASSWORD"`
	DatabaseName string      `default:""          envconfig:"CONNECTIONS_MONGO_NAME"`
	Replset      string      `default:""          envconfig:"CONNECTIONS_MONGO_REPLSET"`
	Timeout      uint        `default:"10"        envconfig:"CONNECTIONS_MONGO_TIMEOUT"`
	SSL          bool        `default:"false"     envconfig:"CONNECTIONS_MONGO_SSL"`
	TLSConfig    *tls.Config `ignored:"true"`
}

// DefaultConfig returns a configuration for a locally hosted, unauthenticated mongo
func DefaultConfig() *Config {
	return &Config{
		Hostnames:    []string{defaultHost},
		Port:         defaultPort,
		DatabaseName: defaultDatabaseName,
	}
}

// ConnectionInfo configures options for establishing a session with a MongoDB cluster.
func ConnectionInfo(cfg *Config) *mgo.DialInfo {
	if len(cfg.Hostnames) == 0 {
		cfg.Hostnames = []string{defaultHost}
	}

	if cfg.DatabaseName == "" {
		cfg.DatabaseName = defaultDatabaseName
	}

	if cfg.Port > 0 {
		for i := range cfg.Hostnames {
			cfg.Hostnames[i] = fmt.Sprintf("%s:%d", cfg.Hostnames[i], cfg.Port)
		}
	}

	if cfg.Timeout == 0 {
		cfg.Timeout = 10
	}

	dialInfo := &mgo.DialInfo{
		Addrs:          cfg.Hostnames,
		Database:       cfg.DatabaseName,
		Username:       cfg.Username,
		Password:       cfg.Password,
		Source:         cfg.AuthDB,
		ReplicaSetName: cfg.Replset,
		Timeout:        time.Second * time.Duration(cfg.Timeout),
	}

	if cfg.SSL {
		dialInfo.DialServer = func(addr *mgo.ServerAddr) (net.Conn, error) {
			return tls.Dial("tcp", addr.String(), cfg.TLSConfig)
		}
	}

	return dialInfo
}

// Connect returns a connection to mongo.
func Connect(cfg *Config) (*mgo.Database, error) {
	log := logger.WithFields(logrus.Fields{
		"package": "mongo",
		"method":  "Connect",
	})

	dialInfo := ConnectionInfo(cfg)
	session, err := mgo.DialWithInfo(dialInfo)
	if err != nil {
		log.WithError(err).Error("Unable to connect to mongo! Have you configured your connection properly?")
		return nil, err
	}

	// Monotonic consistency will start reading from a slave if possible
	session.SetMode(mgo.Monotonic, true)
	return session.DB(cfg.DatabaseName), nil
}

// New allows for custom mongo configuration and custom hashers.
func New(cfg *Config, hashee fosite.Hasher) (*Store, error) {
	log := logger.WithFields(logrus.Fields{
		"package": "mongo",
		"method":  "NewFromConfig",
	})

	database, err := Connect(cfg)
	if err != nil {
		log.WithError(err).Error("Unable to connect to mongo! Are you sure mongo is running on localhost?")
		return nil, err
	}

	if hashee == nil {
		// Initialize default fosite Hasher.
		hashee = &fosite.BCrypt{
			WorkFactor: 10,
		}
	}

	// Build up the mongo endpoints
	mongoCache := &CacheManager{
		DB: database,
	}
	mongoClients := &ClientManager{
		DB:     database,
		Hasher: hashee,
	}
	mongoUsers := &UserManager{
		DB:     database,
		Hasher: hashee,
	}
	mongoRequests := &RequestManager{
		DB: database,

		Cache:   mongoCache,
		Clients: mongoClients,
		Users:   mongoUsers,
	}

	// Init DB collections, indices e.t.c.
	managers := []storage.Configurer{
		mongoCache,
		mongoClients,
		mongoUsers,
		mongoRequests,
	}

	// Use the main DB database to configure the mongo collections on first up.
	mgoSession := database.Session.Copy()
	defer mgoSession.Close()
	ctx := MgoSessionToContext(context.Background(), mgoSession)

	for _, manager := range managers {
		err := manager.Configure(ctx)
		if err != nil {
			log.WithError(err).Error("Unable to configure mongo collections!")
			return nil, err
		}
	}

	return &Store{
		DB:     database,
		Hasher: hashee,
		Store: storage.Store{
			CacheManager:   mongoCache,
			ClientManager:  mongoClients,
			RequestManager: mongoRequests,
			UserManager:    mongoUsers,
		},
	}, nil
}

// NewDefaultStore returns a Store configured with the default mongo
// configuration and default Hasher.
func NewDefaultStore() (*Store, error) {
	cfg := DefaultConfig()
	return New(cfg, nil)
}
<|endoftext|>


<commit_before>// This is the less generic tree, it currently only supports
// integers, keys (unevaluated operators), and variables (unevaluated constants)

package optim

import (
	"../parser"
	"../token"
	"errors"
	"fmt"
	"log"
	"strconv"
)

type ItemType int

const (
	ItemInt ItemType = iota
	ItemVar
	ItemKey
)

type Tree struct {
	Val *Node
	Sub []*Tree
}

type Node struct {
	Typ    ItemType
	Num    int            // if it is an int
	Var    string         // if it is a variable
	Key    token.ItemType // deferred keyword
	Solved bool           // var is either assigned a value or unknown
}

func (tree *Tree) Append(node *Node) *Tree {
	tree.Sub = append(tree.Sub, &Tree{
		Val: node,
	})
	return tree.Sub[len(tree.Sub)-1]
}

func (tree *Tree) Walk(level int) (*Tree, error) {
	if level != 0 {
		if len(tree.Sub) > 0 {
			return tree.Sub[len(tree.Sub)-1].Walk(level - 1)
		} else {
			return nil, errors.New("level nonexistant")
		}
	} else {
		return tree, nil
	}
}

func (tree *Tree) String() string {
	var s string
	if tree.Val != nil {
		s += tree.Val.String()
	}
	if len(tree.Sub) > 0 {
		s += "{"
		for i := 0; i < len(tree.Sub); i++ {
			if i != len(tree.Sub)-1 {
				s += tree.Sub[i].String() + ", "
			} else {
				s += tree.Sub[i].String()
			}
		}
		s += "}"
	}
	return s
}

func (node *Node) String() string {
	switch node.Typ {
	case ItemInt:
		return "\"" + strconv.Itoa(node.Num) + "\""
	case ItemVar:
		if node.Solved {
			return "(" + node.Var + ":" + strconv.Itoa(node.Num) + ")"
		} else {
			return "(" + node.Var + ")"
		}
	case ItemKey:
		return "symb"
	default:
		return "unk"
	}
}

type evals struct {
	Root      *Tree        // root of tree
	Tree      *Tree        // current branch
	ParseRoot *parser.Tree // root of parse tree
	ParseTree *parser.Tree // root of parse tree
	level     int
}

func Eval(tree *parser.Tree) *Tree {
	t := new(Tree)
	e := &evals{
		Root:      t,
		Tree:      t,
		ParseRoot: tree,
		ParseTree: tree,
	}
	e.createTree(e.ParseRoot, e.Root)
	e.evaluate(e.Root)
	return e.Root
}

var lookup = map[token.ItemType]eval{
	token.ItemAdd: evalAdd,
	token.ItemSub: evalSub,
	token.ItemMul: evalMul,
	token.ItemDiv: evalDiv,
}

// Creates a tree from the parse tree
func (e *evals) createTree(t *parser.Tree, tr *Tree) {
	// check fo' nills
	if t.Val == nil {
		if !hasActionChildren(t) {
			return
		}
	} else {
		//fmt.Printf("%d\n", t.Val.Tok.Typ)
	}

	// We can do stuff
	if t.Val == nil && hasActionChildren(t) {
		for i := 0; i < len(t.Sub); i++ {
			e.createTree(t.Sub[i], tr)
		}
	} else if token.Keyword(t.Val.Tok.Typ) {
		tr = tr.Append(&Node{
			Typ: ItemKey,
			Key: t.Val.Tok.Typ,
		})
		for i := 0; i < len(t.Sub); i++ {
			e.createTree(t.Sub[i], tr)
		}
	} else if token.Constant(t.Val.Tok.Typ) {
		num, err := strconv.Atoi(t.Val.Tok.Val)
		if err != nil {
			log.Fatal(err)
		}
		tr.Append(&Node{
			Typ: ItemInt,
			Num: num,
		})
	} else if t.Val.Tok.Typ == token.ItemVariable {
		tr = tr.Append(&Node{
			Typ: ItemVar,
			Var: t.Val.Tok.Val,
		})
		for i := 0; i < len(t.Sub); i++ {
			e.createTree(t.Sub[i], tr)
		}
	}
}

// evaluate does all the maths it can
func (e *evals) evaluate(t *Tree) *Tree {
	// evaluate valueless trees that contain children
	if t.Val == nil {
		for i := 0; i < len(t.Sub); i++ {
			tree := e.evaluate(t.Sub[i])
			if tree != nil {
				t.Sub[i] = tree
			}
		}
		// evaluate keys
	} else if t.Val.Typ == ItemKey {
		// if the key is an assignment key (special case)
		// call variables to compute value
		if t.Val.Key == token.ItemAssign {
			return e.variables(t)
		}
		// there are keys/vars that need to be computed,
		// recurse to compute them
		if hasSomeKeyChildren(t) || hasSomeVarChildren(t) {
			for i := 0; i < len(t.Sub); i++ {
				// recurse on key
				if t.Sub[i].Val.Typ == ItemKey {
					tree := e.evaluate(t.Sub[i])
					if tree != nil {
						t.Sub[i] = tree
					}
					// recurse on var
				} else if t.Sub[i].Val.Typ == ItemVar {
					tree := e.evaluate(t.Sub[i])
					if tree != nil {
						t.Sub[i] = tree
					}
				}
			}
		}
		// At this point there should only be ints,
		// if so, compute
		if hasOnlyIntChildren(t) {
			if val, ok := lookup[t.Val.Key]; ok {
				return val(t)
			} else {
				return nil
			}
		}
		// evaluate variables if we run into any
	} else if t.Val.Typ == ItemVar {
		return e.variables(t)
	}
	return nil
}

type Var struct {
	Solved bool   // tells when Num is 0 and when Num is empty
	Var    string // variable name
	Num    int    // integer value
}

type Variab struct {
	Var []*Var // array of Vars to create a list of variables
}

var variabs Variab // global list of variables

func (v Variab) getName(s string) *Var {
	for i := 0; i < len(v.Var); i++ {
		if v.Var[i].Var == s {
			return v.Var[i]
		}
	}
	return nil
}

// Variables is called when (1) there is an assignment key
// (2) there is a variable
func (e *evals) variables(t *Tree) *Tree {
	// evaluate assignment keys
	if t.Val.Key == token.ItemAssign {
		// itemAssign is the assignment operator, look for variables
		// assignment must have two operators, a variable & some assigned value (can be a key)
		if len(t.Sub) == 2 && t.Sub[0].Val.Typ == ItemVar {
			// evaluate the key
			if t.Sub[1].Val.Typ == ItemKey {
				tree := e.evaluate(t.Sub[1])
				if tree != nil {
					t.Sub[1] = tree
				}
				// evaluate variable
			} else if t.Sub[1].Val.Typ == ItemVar {
				// replace var with var's value node if the var is known
				tree := e.variables(t.Sub[1])
				if tree != nil {
					t.Sub[1] = tree
				}
				return e.variables(t)
			}
			// should be int at this point
			if t.Sub[1].Val.Typ == ItemInt {
				// note new variable
				variabs.Var = append(variabs.Var, &Var{
					Var:    t.Sub[0].Val.Var,
					Solved: true,
					Num:    t.Sub[1].Val.Num,
				})
				// return tree
				return &Tree{
					Val: &Node{
						Typ:    ItemVar,
						Num:    t.Sub[1].Val.Num,
						Var:    t.Sub[0].Val.Var,
						Solved: true,
					},
				}
			}
		} else {
			fmt.Printf("Incorrect usage of ':' operator.")
			return nil
		}
		// evaluate variables
	} else if t.Val.Typ == ItemVar {
		// t.Val.Var is in the list of found variables
		if v := variabs.getName(t.Val.Var); v != nil {
			// return the tree, which will be appended by its parent in its place
			return &Tree{
				Val: &Node{
					Typ:    ItemVar,
					Num:    v.Num,
					Var:    v.Var,
					Solved: v.Solved,
				},
			}
		}
	}
	return nil
}

// evals
type eval func(t *Tree) *Tree

func evalAdd(t *Tree) *Tree {
	n := t.Sub[0].Val.Num
	for i := 1; i < len(t.Sub); i++ {
		n += t.Sub[i].Val.Num
	}
	return &Tree{
		Val: &Node{
			Typ: ItemInt,
			Num: n,
		},
	}
	//fmt.Printf("add %s = %d\n", num, n)
}

func evalSub(t *Tree) *Tree {
	n := t.Sub[0].Val.Num
	for i := 1; i < len(t.Sub); i++ {
		n -= t.Sub[i].Val.Num
	}
	return &Tree{
		Val: &Node{
			Typ: ItemInt,
			Num: n,
		},
	}
	//fmt.Printf("subtract %s = %d\n", num, n)
}

func evalMul(t *Tree) *Tree {
	n := t.Sub[0].Val.Num
	for i := 1; i < len(t.Sub); i++ {
		n *= t.Sub[i].Val.Num
	}
	return &Tree{
		Val: &Node{
			Typ: ItemInt,
			Num: n,
		},
	}
	//fmt.Printf("multiply %s = %d\n", num, n)
}

func evalDiv(t *Tree) *Tree {
	n := t.Sub[0].Val.Num
	for i := 1; i < len(t.Sub); i++ {
		n /= t.Sub[i].Val.Num
	}
	return &Tree{
		Val: &Node{
			Typ: ItemInt,
			Num: n,
		},
	}
	//fmt.Printf("divide %s = %d\n", num, n)
}

func hasActionChildren(tree *parser.Tree) bool {
	if tree.Sub != nil && len(tree.Sub) > 0 {
		for i := 0; i < len(tree.Sub); i++ {
			if token.Keyword(tree.Sub[i].Val.Tok.Typ) {
				return true
			}
		}
	}
	return false
}

func hasOnlyIntChildren(tree *Tree) bool {
	if tree.Sub != nil && len(tree.Sub) > 0 {
		for i := 0; i < len(tree.Sub); i++ {
			if tree.Sub[i].Val.Typ != ItemInt && !(tree.Sub[i].Val.Typ == ItemVar && tree.Sub[i].Val.Solved == true) {
				return false
			}
		}
		return true
	}
	return false
}

func hasSomeKeyChildren(tree *Tree) bool {
	if tree.Sub != nil && len(tree.Sub) > 0 {
		for i := 0; i < len(tree.Sub); i++ {
			if tree.Sub[i].Val.Typ == ItemKey {
				return true
			}
		}
	}
	return false
}

func hasSomeVarChildren(tree *Tree) bool {
	if tree.Sub != nil && len(tree.Sub) > 0 {
		for i := 0; i < len(tree.Sub); i++ {
			if tree.Sub[i].Val.Typ == ItemVar {
				return true
			}
		}
	}
	return false
}

func isAction(tree *parser.Tree) bool {
	if token.Keyword(tree.Val.Tok.Typ) {
		return true
	}
	return false
}

func intify(tree *parser.Tree) ([]int, error) {
	var num []int
	for i := 0; i < len(tree.Sub); i++ {
		switch tree.Sub[i].Val.Tok.Typ {
		case token.ItemNumber:
			n, err := strconv.Atoi(tree.Sub[i].Val.Tok.Val)
			if err != nil {
				return nil, err
			}
			num = append(num, n)
		}
	}
	return num, nil
}
<commit_msg>assign existing var a value works<commit_after>// This is the less generic tree, it currently only supports
// integers, keys (unevaluated operators), and variables (unevaluated constants)

package optim

import (
	"../parser"
	"../token"
	"errors"
	"fmt"
	"log"
	"strconv"
)

type ItemType int

const (
	ItemInt ItemType = iota
	ItemVar
	ItemKey
)

type Tree struct {
	Val *Node
	Sub []*Tree
}

type Node struct {
	Typ    ItemType
	Num    int            // if it is an int
	Var    string         // if it is a variable
	Key    token.ItemType // deferred keyword
	Solved bool           // var is either assigned a value or unknown
}

func (tree *Tree) Append(node *Node) *Tree {
	tree.Sub = append(tree.Sub, &Tree{
		Val: node,
	})
	return tree.Sub[len(tree.Sub)-1]
}

func (tree *Tree) Walk(level int) (*Tree, error) {
	if level != 0 {
		if len(tree.Sub) > 0 {
			return tree.Sub[len(tree.Sub)-1].Walk(level - 1)
		} else {
			return nil, errors.New("level nonexistant")
		}
	} else {
		return tree, nil
	}
}

func (tree *Tree) String() string {
	var s string
	if tree.Val != nil {
		s += tree.Val.String()
	}
	if len(tree.Sub) > 0 {
		s += "{"
		for i := 0; i < len(tree.Sub); i++ {
			if i != len(tree.Sub)-1 {
				s += tree.Sub[i].String() + ", "
			} else {
				s += tree.Sub[i].String()
			}
		}
		s += "}"
	}
	return s
}

func (node *Node) String() string {
	switch node.Typ {
	case ItemInt:
		return "\"" + strconv.Itoa(node.Num) + "\""
	case ItemVar:
		if node.Solved {
			return "(" + node.Var + ":" + strconv.Itoa(node.Num) + ")"
		} else {
			return "(" + node.Var + ")"
		}
	case ItemKey:
		return "symb"
	default:
		return "unk"
	}
}

type evals struct {
	Root      *Tree        // root of tree
	Tree      *Tree        // current branch
	ParseRoot *parser.Tree // root of parse tree
	ParseTree *parser.Tree // root of parse tree
	level     int
}

func Eval(tree *parser.Tree) *Tree {
	t := new(Tree)
	e := &evals{
		Root:      t,
		Tree:      t,
		ParseRoot: tree,
		ParseTree: tree,
	}
	e.createTree(e.ParseRoot, e.Root)
	e.evaluate(e.Root)
	return e.Root
}

var lookup = map[token.ItemType]eval{
	token.ItemAdd: evalAdd,
	token.ItemSub: evalSub,
	token.ItemMul: evalMul,
	token.ItemDiv: evalDiv,
}

// Creates a tree from the parse tree
func (e *evals) createTree(t *parser.Tree, tr *Tree) {
	// check fo' nills
	if t.Val == nil {
		if !hasActionChildren(t) {
			return
		}
	} else {
		//fmt.Printf("%d\n", t.Val.Tok.Typ)
	}

	// We can do stuff
	if t.Val == nil && hasActionChildren(t) {
		for i := 0; i < len(t.Sub); i++ {
			e.createTree(t.Sub[i], tr)
		}
	} else if token.Keyword(t.Val.Tok.Typ) {
		tr = tr.Append(&Node{
			Typ: ItemKey,
			Key: t.Val.Tok.Typ,
		})
		for i := 0; i < len(t.Sub); i++ {
			e.createTree(t.Sub[i], tr)
		}
	} else if token.Constant(t.Val.Tok.Typ) {
		num, err := strconv.Atoi(t.Val.Tok.Val)
		if err != nil {
			log.Fatal(err)
		}
		tr.Append(&Node{
			Typ: ItemInt,
			Num: num,
		})
	} else if t.Val.Tok.Typ == token.ItemVariable {
		tr = tr.Append(&Node{
			Typ: ItemVar,
			Var: t.Val.Tok.Val,
		})
		for i := 0; i < len(t.Sub); i++ {
			e.createTree(t.Sub[i], tr)
		}
	}
}

// evaluate does all the maths it can
func (e *evals) evaluate(t *Tree) *Tree {
	// evaluate valueless trees that contain children
	if t.Val == nil {
		for i := 0; i < len(t.Sub); i++ {
			tree := e.evaluate(t.Sub[i])
			if tree != nil {
				t.Sub[i] = tree
			}
		}
		// evaluate keys
	} else if t.Val.Typ == ItemKey {
		// if the key is an assignment key (special case)
		// call variables to compute value
		if t.Val.Key == token.ItemAssign {
			return e.variables(t)
		}
		// there are keys/vars that need to be computed,
		// recurse to compute them
		if hasSomeKeyChildren(t) || hasSomeVarChildren(t) {
			for i := 0; i < len(t.Sub); i++ {
				// recurse on key
				if t.Sub[i].Val.Typ == ItemKey {
					tree := e.evaluate(t.Sub[i])
					if tree != nil {
						t.Sub[i] = tree
					}
					// recurse on var
				} else if t.Sub[i].Val.Typ == ItemVar {
					tree := e.evaluate(t.Sub[i])
					if tree != nil {
						t.Sub[i] = tree
					}
				}
			}
		}
		// At this point there should only be ints,
		// if so, compute
		if hasOnlyIntChildren(t) {
			if val, ok := lookup[t.Val.Key]; ok {
				return val(t)
			} else {
				return nil
			}
		}
		// evaluate variables if we run into any
	} else if t.Val.Typ == ItemVar {
		return e.variables(t)
	}
	return nil
}

type Var struct {
	Solved bool   // tells when Num is 0 and when Num is empty
	Var    string // variable name
	Num    int    // integer value
}

type Variab struct {
	Var []*Var // array of Vars to create a list of variables
}

var variabs Variab // global list of variables

func (v Variab) getName(s string) *Var {
	for i := 0; i < len(v.Var); i++ {
		if v.Var[i].Var == s {
			return v.Var[i]
		}
	}
	return nil
}

// Variables is called when (1) there is an assignment key
// (2) there is a variable
func (e *evals) variables(t *Tree) *Tree {
	// evaluate assignment keys
	if t.Val.Key == token.ItemAssign {
		// itemAssign is the assignment operator, look for variables
		// assignment must have two operators, a variable & some assigned value (can be a key)
		if len(t.Sub) == 2 && t.Sub[0].Val.Typ == ItemVar {
			// evaluate the key
			if t.Sub[1].Val.Typ == ItemKey {
				tree := e.evaluate(t.Sub[1])
				if tree != nil {
					t.Sub[1] = tree
				}
				// evaluate variable
			} else if t.Sub[1].Val.Typ == ItemVar {
				// replace var with var's value node if the var is known
				tree := e.variables(t.Sub[1])
				if tree != nil {
					t.Sub[1] = tree
				}
				return e.variables(t)
			}
			// should be int at this point
			if t.Sub[1].Val.Typ == ItemInt {
				// check if just reassigning existing variable
				if v := variabs.getName(t.Sub[0].Val.Var); v != nil {
					v.Num = t.Sub[1].Val.Num
					// new variable creation
				} else {
					variabs.Var = append(variabs.Var, &Var{
						Var:    t.Sub[0].Val.Var,
						Solved: true,
						Num:    t.Sub[1].Val.Num,
					})
				}
				// return tree
				return &Tree{
					Val: &Node{
						Typ:    ItemVar,
						Num:    t.Sub[1].Val.Num,
						Var:    t.Sub[0].Val.Var,
						Solved: true,
					},
				}
			}
		} else {
			fmt.Printf("Incorrect usage of ':' operator.")
			return nil
		}
		// evaluate variables
	} else if t.Val.Typ == ItemVar {
		// t.Val.Var is in the list of found variables
		if v := variabs.getName(t.Val.Var); v != nil {
			// return the tree, which will be appended by its parent in its place
			return &Tree{
				Val: &Node{
					Typ:    ItemVar,
					Num:    v.Num,
					Var:    v.Var,
					Solved: v.Solved,
				},
			}
		}
	}
	return nil
}

// evals
type eval func(t *Tree) *Tree

func evalAdd(t *Tree) *Tree {
	n := t.Sub[0].Val.Num
	for i := 1; i < len(t.Sub); i++ {
		n += t.Sub[i].Val.Num
	}
	return &Tree{
		Val: &Node{
			Typ: ItemInt,
			Num: n,
		},
	}
	//fmt.Printf("add %s = %d\n", num, n)
}

func evalSub(t *Tree) *Tree {
	n := t.Sub[0].Val.Num
	for i := 1; i < len(t.Sub); i++ {
		n -= t.Sub[i].Val.Num
	}
	return &Tree{
		Val: &Node{
			Typ: ItemInt,
			Num: n,
		},
	}
	//fmt.Printf("subtract %s = %d\n", num, n)
}

func evalMul(t *Tree) *Tree {
	n := t.Sub[0].Val.Num
	for i := 1; i < len(t.Sub); i++ {
		n *= t.Sub[i].Val.Num
	}
	return &Tree{
		Val: &Node{
			Typ: ItemInt,
			Num: n,
		},
	}
	//fmt.Printf("multiply %s = %d\n", num, n)
}

func evalDiv(t *Tree) *Tree {
	n := t.Sub[0].Val.Num
	for i := 1; i < len(t.Sub); i++ {
		n /= t.Sub[i].Val.Num
	}
	return &Tree{
		Val: &Node{
			Typ: ItemInt,
			Num: n,
		},
	}
	//fmt.Printf("divide %s = %d\n", num, n)
}

func hasActionChildren(tree *parser.Tree) bool {
	if tree.Sub != nil && len(tree.Sub) > 0 {
		for i := 0; i < len(tree.Sub); i++ {
			if token.Keyword(tree.Sub[i].Val.Tok.Typ) {
				return true
			}
		}
	}
	return false
}

func hasOnlyIntChildren(tree *Tree) bool {
	if tree.Sub != nil && len(tree.Sub) > 0 {
		for i := 0; i < len(tree.Sub); i++ {
			if tree.Sub[i].Val.Typ != ItemInt && !(tree.Sub[i].Val.Typ == ItemVar && tree.Sub[i].Val.Solved == true) {
				return false
			}
		}
		return true
	}
	return false
}

func hasSomeKeyChildren(tree *Tree) bool {
	if tree.Sub != nil && len(tree.Sub) > 0 {
		for i := 0; i < len(tree.Sub); i++ {
			if tree.Sub[i].Val.Typ == ItemKey {
				return true
			}
		}
	}
	return false
}

func hasSomeVarChildren(tree *Tree) bool {
	if tree.Sub != nil && len(tree.Sub) > 0 {
		for i := 0; i < len(tree.Sub); i++ {
			if tree.Sub[i].Val.Typ == ItemVar {
				return true
			}
		}
	}
	return false
}

func isAction(tree *parser.Tree) bool {
	if token.Keyword(tree.Val.Tok.Typ) {
		return true
	}
	return false
}

func intify(tree *parser.Tree) ([]int, error) {
	var num []int
	for i := 0; i < len(tree.Sub); i++ {
		switch tree.Sub[i].Val.Tok.Typ {
		case token.ItemNumber:
			n, err := strconv.Atoi(tree.Sub[i].Val.Tok.Val)
			if err != nil {
				return nil, err
			}
			num = append(num, n)
		}
	}
	return num, nil
}
<|endoftext|>


<commit_before>/*
 * Copyright 2017 Dgraph Labs, Inc. and Contributors
 *
 * 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 badger

import (
	"context"
	"fmt"
	"os"
	"sync"
	"time"

	"golang.org/x/net/trace"

	"github.com/dgraph-io/badger/skl"
	"github.com/dgraph-io/badger/table"
	"github.com/dgraph-io/badger/y"
)

var (
	Head = []byte("/head/") // For storing value offset for replay.
)

// Options are params for creating DB object.
type Options struct {
	Dir                     string
	NumLevelZeroTables      int   // Maximum number of Level 0 tables before we start compacting.
	NumLevelZeroTablesStall int   // If we hit this number of Level 0 tables, we will stall until level 0 is compacted away.
	LevelOneSize            int64 // Maximum total size for Level 1.
	MaxLevels               int   // Maximum number of levels of compaction. May be made variable later.
	MaxTableSize            int64 // Each table (or file) is at most this size.
	MemtableSlack           int64 // Arena has to be slightly bigger than MaxTableSize.
	LevelSizeMultiplier     int
	ValueThreshold          int // If value size >= this threshold, we store value offsets in tables.
	Verbose                 bool
	DoNotCompact            bool
	MapTablesTo             int
	NumMemtables            int
	ValueGCThreshold        float64
	SyncWrites              bool
}

var DefaultOptions = Options{
	Dir:                     "/tmp",
	NumLevelZeroTables:      5,
	NumLevelZeroTablesStall: 10,
	LevelOneSize:            256 << 20,
	MaxLevels:               7,
	MaxTableSize:            64 << 20,
	LevelSizeMultiplier:     10,
	ValueThreshold:          20,
	Verbose:                 true,
	DoNotCompact:            false, // Only for testing.
	MapTablesTo:             table.MemoryMap,
	NumMemtables:            5,
	MemtableSlack:           10 << 20,
	ValueGCThreshold:        0.8, // Set to zero to not run GC.
	SyncWrites:              true,
}

type KV struct {
	sync.RWMutex // Guards imm, mem.

	closer    *y.Closer
	elog      trace.EventLog
	mt        *skl.Skiplist
	imm       []*skl.Skiplist // Add here only AFTER pushing to flushChan.
	opt       Options
	lc        *levelsController
	vlog      valueLog
	vptr      valuePointer
	arenaPool *skl.ArenaPool
	writeCh   chan *block
	flushChan chan flushTask // For flushing memtables.
}

// NewKV returns a new KV object. Compact levels are created as well.
func NewKV(opt *Options) *KV {
	y.AssertTrue(len(opt.Dir) > 0)
	out := &KV{
		imm:       make([]*skl.Skiplist, 0, opt.NumMemtables),
		flushChan: make(chan flushTask, opt.NumMemtables),
		writeCh:   make(chan *block, 1000),
		opt:       *opt, // Make a copy.
		arenaPool: skl.NewArenaPool(opt.MaxTableSize+opt.MemtableSlack, opt.NumMemtables+5),
		closer:    y.NewCloser(0),
		elog:      trace.NewEventLog("badger", "kv"),
	}
	out.mt = skl.NewSkiplist(out.arenaPool)

	// newLevelsController potentially loads files in directory.
	out.lc = newLevelsController(out)
	out.lc.startCompact()

	out.closer.Register()
	go out.flushMemtable() // Need levels controller to be up.

	out.vlog.Open(out, opt)

	val := out.Get(context.Background(), Head)
	var vptr valuePointer
	if len(val) > 0 {
		vptr.Decode(val)
	}

	first := true
	fn := func(e Entry) bool {
		if first {
			fmt.Printf("key=%s\n", e.Key)
		}
		first = false
		nk := make([]byte, len(e.Key))
		copy(nk, e.Key)
		nv := make([]byte, len(e.Value))
		copy(nv, e.Value)
		out.mt.Put(nk, nv, e.Meta)
		return true
	}
	out.vlog.Replay(vptr, fn)

	out.closer.Register()
	go out.doWrites()
	return out
}

// Close closes a KV.
func (s *KV) Close() {
	if s.opt.Verbose {
		y.Printf("Closing database\n")
	}

	s.vlog.Close() // Close value log first, because it still needs KV to be active.
	s.elog.Printf("Closing database")
	s.closer.Signal()

	// TODO(ganesh): Block writes to mem.
	if s.mt.Size() > 0 {
		if s.opt.Verbose {
			y.Printf("Flushing memtable\n")
		}
		s.Lock()
		y.AssertTrue(s.mt != nil)
		s.flushChan <- flushTask{s.mt, s.vptr}
		s.imm = append(s.imm, s.mt) // Flusher will attempt to remove this from s.imm.
		s.mt = nil                  // Will segfault if we try writing!
		s.Unlock()
	}
	s.flushChan <- flushTask{nil, valuePointer{}} // Tell flusher to quit.
	if s.opt.Verbose {
		y.Printf("Memtable flushed\n")
	}

	s.lc.close()
	s.elog.Printf("Waiting for closer")
	s.closer.Wait()
	s.elog.Finish()
}

func (s *KV) getMemTables() []*skl.Skiplist {
	s.RLock()
	defer s.RUnlock()
	tables := make([]*skl.Skiplist, len(s.imm)+1)
	tables[0] = s.mt
	last := len(s.imm) - 1
	for i := range s.imm {
		tables[i+1] = s.imm[last-i]
	}
	return tables
}

func (s *KV) decodeValue(ctx context.Context, val []byte, meta byte) []byte {
	if (meta & BitDelete) != 0 {
		// Tombstone encountered.
		return nil
	}
	if (meta & BitValuePointer) == 0 {
		y.Trace(ctx, "Value stored with key")
		return val
	}

	var vp valuePointer
	vp.Decode(val)
	entry, err := s.vlog.Read(ctx, vp)
	y.Checkf(err, "Unable to read from value log: %+v", vp)

	if (entry.Meta & BitDelete) == 0 { // Not tombstone.
		return entry.Value
	}
	return []byte{}
}

// getValueHelper returns the value in memtable or disk for given key.
// Note that value will include meta byte.
func (s *KV) get(ctx context.Context, key []byte) ([]byte, byte) {
	y.Trace(ctx, "Retrieving key from memtables")
	tables := s.getMemTables() // Lock should be released.
	for i := 0; i < len(tables); i++ {
		v, meta := tables[i].Get(key)
		if meta != 0 || v != nil {
			return v, meta
		}
	}
	y.Trace(ctx, "Not found in memtables. Getting from levels")
	return s.lc.get(ctx, key)
}

// Get looks for key and returns value. If not found, return nil.
func (s *KV) Get(ctx context.Context, key []byte) []byte {
	val, meta := s.get(ctx, key)
	return s.decodeValue(ctx, val, meta)
}

func (s *KV) updateOffset(ptrs []valuePointer) {
	ptr := ptrs[len(ptrs)-1]

	s.Lock()
	defer s.Unlock()
	if s.vptr.Fid < ptr.Fid {
		s.vptr = ptr
	} else if s.vptr.Offset < ptr.Offset {
		s.vptr = ptr
	}
}

var blockPool = sync.Pool{
	New: func() interface{} {
		return new(block)
	},
}

func (s *KV) writeToLSM(b *block) {
	var offsetBuf [16]byte
	y.AssertTrue(len(b.Ptrs) == len(b.Entries))
	for i, entry := range b.Entries {
		if len(entry.Value) < s.opt.ValueThreshold { // Will include deletion / tombstone case.
			s.mt.Put(entry.Key, entry.Value, entry.Meta)
		} else {
			s.mt.Put(entry.Key, b.Ptrs[i].Encode(offsetBuf[:]), entry.Meta|BitValuePointer)
		}
	}
}

func (s *KV) writeBlocks(blocks []*block) {
	if len(blocks) == 0 {
		return
	}
	s.elog.Printf("writeBlocks called")

	for !s.hasRoomForWrite() {
		s.elog.Printf("Making room for writes")
		// We need to poll a bit because both hasRoomForWrite and the flusher need access to s.imm.
		// When flushChan is full and you are blocked there, and the flusher is trying to update s.imm,
		// you will get a deadlock.
		time.Sleep(10 * time.Millisecond)
	}

	s.elog.Printf("Writing to value log")
	s.vlog.Write(blocks)
	s.elog.Printf("Writing to memtable")

	for i, b := range blocks {
		s.writeToLSM(b)
		s.elog.Printf("Wrote %d entries from block %d", len(b.Entries), i)
	}
	lastb := len(blocks) - 1
	s.updateOffset(blocks[lastb].Ptrs)
	for _, b := range blocks {
		b.Wg.Done()
	}
}

func (s *KV) doWrites() {
	defer s.closer.Done()

	blocks := make([]*block, 0, 10)
	for {
		select {
		case b := <-s.writeCh:
			blocks = append(blocks, b)

		case <-s.closer.HasBeenClosed():
			close(s.writeCh)

			for b := range s.writeCh { // Flush the channel.
				blocks = append(blocks, b)
			}
			s.writeBlocks(blocks)
			return

		default:
			if len(blocks) == 0 {
				time.Sleep(time.Millisecond)
				break
			}
			s.writeBlocks(blocks)
			blocks = blocks[:0]
		}
	}
}

// Write applies a list of value.Entry to our memtable.
func (s *KV) Write(ctx context.Context, entries []*Entry) error {
	b := blockPool.Get().(*block)
	defer blockPool.Put(b)

	b.Entries = entries
	b.Wg = sync.WaitGroup{}
	b.Wg.Add(1)
	s.writeCh <- b
	b.Wg.Wait()

	return nil
}

// Put puts a key-val pair.
func (s *KV) Put(ctx context.Context, key []byte, val []byte) error {
	e := &Entry{
		Key:   key,
		Value: val,
	}

	return s.Write(ctx, []*Entry{e})
}

// Delete deletes a key.
func (s *KV) Delete(ctx context.Context, key []byte) error {
	e := &Entry{
		Key:  key,
		Meta: BitDelete,
	}
	return s.Write(ctx, []*Entry{e})
}

func (s *KV) hasRoomForWrite() bool {
	s.Lock()
	defer s.Unlock()
	if s.mt.Size() < s.opt.MaxTableSize {
		return true
	}

	y.AssertTrue(s.mt != nil) // A nil mt indicates that KV is being closed.
	select {
	case s.flushChan <- flushTask{s.mt, s.vptr}:
		if s.opt.Verbose {
			y.Printf("Flushing memtable, mt.size=%d size of flushChan: %d\n",
				s.mt.Size(), len(s.flushChan))
		}
		// We manage to push this task. Let's modify imm.
		s.imm = append(s.imm, s.mt)
		s.mt = skl.NewSkiplist(s.arenaPool)
		// New memtable is empty. We certainly have room.
		return true
	default:
		// We need to do this to unlock and allow the flusher to modify imm.
		return false
	}
}

// WriteLevel0Table flushes memtable. It drops deleteValues.
func writeLevel0Table(s *skl.Skiplist, f *os.File) error {
	iter := s.NewIterator()
	defer iter.Close()
	b := table.NewTableBuilder()
	defer b.Close()
	for iter.SeekToFirst(); iter.Valid(); iter.Next() {
		val, meta := iter.Value()
		if err := b.Add(iter.Key(), val, meta); err != nil {
			return err
		}
	}
	var buf [2]byte // Level 0. Leave it initialized as 0.
	_, err := f.Write(b.Finish(buf[:]))
	return err
}

type flushTask struct {
	mt   *skl.Skiplist
	vptr valuePointer
}

func (s *KV) flushMemtable() {
	defer s.closer.Done()

	for {
		select {
		case ft := <-s.flushChan:
			if ft.mt == nil {
				return
			}

			if ft.vptr.Fid > 0 || ft.vptr.Offset > 0 {
				if s.opt.Verbose {
					fmt.Printf("Storing offset: %+v\n", ft.vptr)
				}
				offset := make([]byte, 16)
				s.vptr.Encode(offset)
				ft.mt.Put(Head, offset, 0)
			}
			fileID, _ := s.lc.reserveFileIDs(1)
			fd, err := y.OpenSyncedFile(table.NewFilename(fileID, s.opt.Dir), true)
			y.Check(err)
			y.Check(writeLevel0Table(ft.mt, fd))

			tbl, err := table.OpenTable(fd, s.opt.MapTablesTo)
			defer tbl.DecrRef()

			y.Check(err)
			s.lc.addLevel0Table(tbl) // This will incrRef again.

			// Update s.imm. Need a lock.
			s.Lock()
			y.AssertTrue(ft.mt == s.imm[0]) //For now, single threaded.
			s.imm = s.imm[1:]
			ft.mt.DecrRef() // Return memory.
			s.Unlock()
		}
	}
}
<commit_msg>Check space usage for memtable after every block<commit_after>/*
 * Copyright 2017 Dgraph Labs, Inc. and Contributors
 *
 * 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 badger

import (
	"context"
	"fmt"
	"os"
	"sync"
	"time"

	"golang.org/x/net/trace"

	"github.com/dgraph-io/badger/skl"
	"github.com/dgraph-io/badger/table"
	"github.com/dgraph-io/badger/y"
)

var (
	Head = []byte("/head/") // For storing value offset for replay.
)

// Options are params for creating DB object.
type Options struct {
	Dir                     string
	NumLevelZeroTables      int   // Maximum number of Level 0 tables before we start compacting.
	NumLevelZeroTablesStall int   // If we hit this number of Level 0 tables, we will stall until level 0 is compacted away.
	LevelOneSize            int64 // Maximum total size for Level 1.
	MaxLevels               int   // Maximum number of levels of compaction. May be made variable later.
	MaxTableSize            int64 // Each table (or file) is at most this size.
	MemtableSlack           int64 // Arena has to be slightly bigger than MaxTableSize.
	LevelSizeMultiplier     int
	ValueThreshold          int // If value size >= this threshold, we store value offsets in tables.
	Verbose                 bool
	DoNotCompact            bool
	MapTablesTo             int
	NumMemtables            int
	ValueGCThreshold        float64
	SyncWrites              bool
}

var DefaultOptions = Options{
	Dir:                     "/tmp",
	NumLevelZeroTables:      5,
	NumLevelZeroTablesStall: 10,
	LevelOneSize:            256 << 20,
	MaxLevels:               7,
	MaxTableSize:            64 << 20,
	LevelSizeMultiplier:     10,
	ValueThreshold:          20,
	Verbose:                 true,
	DoNotCompact:            false, // Only for testing.
	MapTablesTo:             table.MemoryMap,
	NumMemtables:            5,
	MemtableSlack:           10 << 20,
	ValueGCThreshold:        0.8, // Set to zero to not run GC.
	SyncWrites:              true,
}

type KV struct {
	sync.RWMutex // Guards imm, mem.

	closer    *y.Closer
	elog      trace.EventLog
	mt        *skl.Skiplist
	imm       []*skl.Skiplist // Add here only AFTER pushing to flushChan.
	opt       Options
	lc        *levelsController
	vlog      valueLog
	vptr      valuePointer
	arenaPool *skl.ArenaPool
	writeCh   chan *block
	flushChan chan flushTask // For flushing memtables.
}

// NewKV returns a new KV object. Compact levels are created as well.
func NewKV(opt *Options) *KV {
	y.AssertTrue(len(opt.Dir) > 0)
	out := &KV{
		imm:       make([]*skl.Skiplist, 0, opt.NumMemtables),
		flushChan: make(chan flushTask, opt.NumMemtables),
		writeCh:   make(chan *block, 1000),
		opt:       *opt, // Make a copy.
		arenaPool: skl.NewArenaPool(opt.MaxTableSize+opt.MemtableSlack, opt.NumMemtables+5),
		closer:    y.NewCloser(0),
		elog:      trace.NewEventLog("badger", "kv"),
	}
	out.mt = skl.NewSkiplist(out.arenaPool)

	// newLevelsController potentially loads files in directory.
	out.lc = newLevelsController(out)
	out.lc.startCompact()

	out.closer.Register()
	go out.flushMemtable() // Need levels controller to be up.

	out.vlog.Open(out, opt)

	val := out.Get(context.Background(), Head)
	var vptr valuePointer
	if len(val) > 0 {
		vptr.Decode(val)
	}

	first := true
	fn := func(e Entry) bool {
		if first {
			fmt.Printf("key=%s\n", e.Key)
		}
		first = false
		nk := make([]byte, len(e.Key))
		copy(nk, e.Key)
		nv := make([]byte, len(e.Value))
		copy(nv, e.Value)
		out.mt.Put(nk, nv, e.Meta)
		return true
	}
	out.vlog.Replay(vptr, fn)

	out.closer.Register()
	go out.doWrites()
	return out
}

// Close closes a KV.
func (s *KV) Close() {
	if s.opt.Verbose {
		y.Printf("Closing database\n")
	}

	s.vlog.Close() // Close value log first, because it still needs KV to be active.
	s.elog.Printf("Closing database")
	s.closer.Signal()

	// TODO(ganesh): Block writes to mem.
	if s.mt.Size() > 0 {
		if s.opt.Verbose {
			y.Printf("Flushing memtable\n")
		}
		s.Lock()
		y.AssertTrue(s.mt != nil)
		s.flushChan <- flushTask{s.mt, s.vptr}
		s.imm = append(s.imm, s.mt) // Flusher will attempt to remove this from s.imm.
		s.mt = nil                  // Will segfault if we try writing!
		s.Unlock()
	}
	s.flushChan <- flushTask{nil, valuePointer{}} // Tell flusher to quit.
	if s.opt.Verbose {
		y.Printf("Memtable flushed\n")
	}

	s.lc.close()
	s.elog.Printf("Waiting for closer")
	s.closer.Wait()
	s.elog.Finish()
}

func (s *KV) getMemTables() []*skl.Skiplist {
	s.RLock()
	defer s.RUnlock()
	tables := make([]*skl.Skiplist, len(s.imm)+1)
	tables[0] = s.mt
	last := len(s.imm) - 1
	for i := range s.imm {
		tables[i+1] = s.imm[last-i]
	}
	return tables
}

func (s *KV) decodeValue(ctx context.Context, val []byte, meta byte) []byte {
	if (meta & BitDelete) != 0 {
		// Tombstone encountered.
		return nil
	}
	if (meta & BitValuePointer) == 0 {
		y.Trace(ctx, "Value stored with key")
		return val
	}

	var vp valuePointer
	vp.Decode(val)
	entry, err := s.vlog.Read(ctx, vp)
	y.Checkf(err, "Unable to read from value log: %+v", vp)

	if (entry.Meta & BitDelete) == 0 { // Not tombstone.
		return entry.Value
	}
	return []byte{}
}

// getValueHelper returns the value in memtable or disk for given key.
// Note that value will include meta byte.
func (s *KV) get(ctx context.Context, key []byte) ([]byte, byte) {
	y.Trace(ctx, "Retrieving key from memtables")
	tables := s.getMemTables() // Lock should be released.
	for i := 0; i < len(tables); i++ {
		v, meta := tables[i].Get(key)
		if meta != 0 || v != nil {
			return v, meta
		}
	}
	y.Trace(ctx, "Not found in memtables. Getting from levels")
	return s.lc.get(ctx, key)
}

// Get looks for key and returns value. If not found, return nil.
func (s *KV) Get(ctx context.Context, key []byte) []byte {
	val, meta := s.get(ctx, key)
	return s.decodeValue(ctx, val, meta)
}

func (s *KV) updateOffset(ptrs []valuePointer) {
	ptr := ptrs[len(ptrs)-1]

	s.Lock()
	defer s.Unlock()
	if s.vptr.Fid < ptr.Fid {
		s.vptr = ptr
	} else if s.vptr.Offset < ptr.Offset {
		s.vptr = ptr
	}
}

var blockPool = sync.Pool{
	New: func() interface{} {
		return new(block)
	},
}

func (s *KV) writeToLSM(b *block) {
	var offsetBuf [16]byte
	y.AssertTrue(len(b.Ptrs) == len(b.Entries))
	for i, entry := range b.Entries {
		if len(entry.Value) < s.opt.ValueThreshold { // Will include deletion / tombstone case.
			s.mt.Put(entry.Key, entry.Value, entry.Meta)
		} else {
			s.mt.Put(entry.Key, b.Ptrs[i].Encode(offsetBuf[:]), entry.Meta|BitValuePointer)
		}
	}
}

func (s *KV) writeBlocks(blocks []*block) {
	if len(blocks) == 0 {
		return
	}
	s.elog.Printf("writeBlocks called")

	s.elog.Printf("Writing to value log")
	s.vlog.Write(blocks)

	s.elog.Printf("Writing to memtable")
	for i, b := range blocks {
		for !s.hasRoomForWrite() {
			s.elog.Printf("Making room for writes")
			// We need to poll a bit because both hasRoomForWrite and the flusher need access to s.imm.
			// When flushChan is full and you are blocked there, and the flusher is trying to update s.imm,
			// you will get a deadlock.
			time.Sleep(10 * time.Millisecond)
		}
		s.writeToLSM(b)
		s.elog.Printf("Wrote %d entries from block %d", len(b.Entries), i)
		s.updateOffset(b.Ptrs)
		b.Wg.Done()
	}
}

func (s *KV) doWrites() {
	defer s.closer.Done()

	blocks := make([]*block, 0, 10)
	for {
		select {
		case b := <-s.writeCh:
			blocks = append(blocks, b)

		case <-s.closer.HasBeenClosed():
			close(s.writeCh)

			for b := range s.writeCh { // Flush the channel.
				blocks = append(blocks, b)
			}
			s.writeBlocks(blocks)
			return

		default:
			if len(blocks) == 0 {
				time.Sleep(time.Millisecond)
				break
			}
			s.writeBlocks(blocks)
			blocks = blocks[:0]
		}
	}
}

// Write applies a list of value.Entry to our memtable.
func (s *KV) Write(ctx context.Context, entries []*Entry) error {
	b := blockPool.Get().(*block)
	defer blockPool.Put(b)

	b.Entries = entries
	b.Wg = sync.WaitGroup{}
	b.Wg.Add(1)
	s.writeCh <- b
	b.Wg.Wait()

	return nil
}

// Put puts a key-val pair.
func (s *KV) Put(ctx context.Context, key []byte, val []byte) error {
	e := &Entry{
		Key:   key,
		Value: val,
	}

	return s.Write(ctx, []*Entry{e})
}

// Delete deletes a key.
func (s *KV) Delete(ctx context.Context, key []byte) error {
	e := &Entry{
		Key:  key,
		Meta: BitDelete,
	}
	return s.Write(ctx, []*Entry{e})
}

func (s *KV) hasRoomForWrite() bool {
	s.Lock()
	defer s.Unlock()
	if s.mt.Size() < s.opt.MaxTableSize {
		return true
	}

	y.AssertTrue(s.mt != nil) // A nil mt indicates that KV is being closed.
	select {
	case s.flushChan <- flushTask{s.mt, s.vptr}:
		if s.opt.Verbose {
			y.Printf("Flushing memtable, mt.size=%d size of flushChan: %d\n",
				s.mt.Size(), len(s.flushChan))
		}
		// We manage to push this task. Let's modify imm.
		s.imm = append(s.imm, s.mt)
		s.mt = skl.NewSkiplist(s.arenaPool)
		// New memtable is empty. We certainly have room.
		return true
	default:
		// We need to do this to unlock and allow the flusher to modify imm.
		return false
	}
}

// WriteLevel0Table flushes memtable. It drops deleteValues.
func writeLevel0Table(s *skl.Skiplist, f *os.File) error {
	iter := s.NewIterator()
	defer iter.Close()
	b := table.NewTableBuilder()
	defer b.Close()
	for iter.SeekToFirst(); iter.Valid(); iter.Next() {
		val, meta := iter.Value()
		if err := b.Add(iter.Key(), val, meta); err != nil {
			return err
		}
	}
	var buf [2]byte // Level 0. Leave it initialized as 0.
	_, err := f.Write(b.Finish(buf[:]))
	return err
}

type flushTask struct {
	mt   *skl.Skiplist
	vptr valuePointer
}

func (s *KV) flushMemtable() {
	defer s.closer.Done()

	for {
		select {
		case ft := <-s.flushChan:
			if ft.mt == nil {
				return
			}

			if ft.vptr.Fid > 0 || ft.vptr.Offset > 0 {
				if s.opt.Verbose {
					fmt.Printf("Storing offset: %+v\n", ft.vptr)
				}
				offset := make([]byte, 16)
				s.vptr.Encode(offset)
				ft.mt.Put(Head, offset, 0)
			}
			fileID, _ := s.lc.reserveFileIDs(1)
			fd, err := y.OpenSyncedFile(table.NewFilename(fileID, s.opt.Dir), true)
			y.Check(err)
			y.Check(writeLevel0Table(ft.mt, fd))

			tbl, err := table.OpenTable(fd, s.opt.MapTablesTo)
			defer tbl.DecrRef()

			y.Check(err)
			s.lc.addLevel0Table(tbl) // This will incrRef again.

			// Update s.imm. Need a lock.
			s.Lock()
			y.AssertTrue(ft.mt == s.imm[0]) //For now, single threaded.
			s.imm = s.imm[1:]
			ft.mt.DecrRef() // Return memory.
			s.Unlock()
		}
	}
}
<|endoftext|>