<commit_before>package agentendpoint

import (
	"context"
	"fmt"
	"reflect"
	"strings"
	"time"

	"github.com/GoogleCloudPlatform/osconfig/agentconfig"
	"github.com/GoogleCloudPlatform/osconfig/attributes"
	"github.com/GoogleCloudPlatform/osconfig/clog"
	"github.com/GoogleCloudPlatform/osconfig/inventory"
	"github.com/GoogleCloudPlatform/osconfig/packages"
	"github.com/GoogleCloudPlatform/osconfig/retryutil"
	"google.golang.org/protobuf/types/known/timestamppb"

	agentendpointpb "google.golang.org/genproto/googleapis/cloud/osconfig/agentendpoint/v1"
	datepb "google.golang.org/genproto/googleapis/type/date"
)

const (
	inventoryURL = agentconfig.ReportURL + "/guestInventory"
)

// ReportInventory writes inventory to guest attributes and reports it to agent endpoint.
func (c *Client) ReportInventory(ctx context.Context) {
	state := inventory.Get(ctx)
	if !agentconfig.DisableInventoryWrite() {
		write(ctx, state, inventoryURL)
	}
	c.report(ctx, state)
}

func write(ctx context.Context, state *inventory.InstanceInventory, url string) {
	clog.Debugf(ctx, "Writing instance inventory to guest attributes.")

	e := reflect.ValueOf(state).Elem()
	t := e.Type()
	for i := 0; i < e.NumField(); i++ {
		f := e.Field(i)
		u := fmt.Sprintf("%s/%s", url, t.Field(i).Name)
		switch f.Kind() {
		case reflect.String:
			clog.Debugf(ctx, "postAttribute %s: %+v", u, f)
			if err := attributes.PostAttribute(u, strings.NewReader(f.String())); err != nil {
				clog.Errorf(ctx, "postAttribute error: %v", err)
			}
		case reflect.Ptr:
			switch reflect.Indirect(f).Kind() {
			case reflect.Struct:
				clog.Debugf(ctx, "postAttributeCompressed %s: %+v", u, f)
				if err := attributes.PostAttributeCompressed(u, f.Interface()); err != nil {
					clog.Errorf(ctx, "postAttributeCompressed error: %v", err)
				}
			}
		}
	}
}

func (c *Client) report(ctx context.Context, state *inventory.InstanceInventory) {
	clog.Debugf(ctx, "Reporting instance inventory to agent endpoint.")
	inventory := formatInventory(ctx, state)

	reportFull := false
	var res *agentendpointpb.ReportInventoryResponse
	var err error
	f := func() error {
		res, err = c.reportInventory(ctx, inventory, reportFull)
		if err != nil {
			return err
		}
		return nil
	}

	if err = retryutil.RetryAPICall(ctx, apiRetrySec*time.Second, "ReportInventory", f); err != nil {
		clog.Errorf(ctx, "Error reporting inventory checksum: %v", err)
		return
	}

	if res.GetReportFullInventory() {
		reportFull = true
		if err = retryutil.RetryAPICall(ctx, apiRetrySec*time.Second, "ReportInventory", f); err != nil {
			clog.Errorf(ctx, "Error reporting full inventory: %v", err)
			return
		}
	}
}

func formatInventory(ctx context.Context, state *inventory.InstanceInventory) *agentendpointpb.Inventory {
	osInfo := &agentendpointpb.Inventory_OsInfo{
		Hostname:             state.Hostname,
		LongName:             state.LongName,
		ShortName:            state.ShortName,
		Version:              state.Version,
		Architecture:         state.Architecture,
		KernelVersion:        state.KernelVersion,
		KernelRelease:        state.KernelRelease,
		OsconfigAgentVersion: state.OSConfigAgentVersion,
	}
	installedPackages := formatPackages(ctx, state.InstalledPackages, state.ShortName)
	availablePackages := formatPackages(ctx, state.PackageUpdates, state.ShortName)

	return &agentendpointpb.Inventory{OsInfo: osInfo, InstalledPackages: installedPackages, AvailablePackages: availablePackages}
}

func formatPackages(ctx context.Context, pkgs *packages.Packages, shortName string) []*agentendpointpb.Inventory_SoftwarePackage {
	var softwarePackages []*agentendpointpb.Inventory_SoftwarePackage
	if pkgs == nil {
		return softwarePackages
	}
	if pkgs.Apt != nil {
		temp := make([]*agentendpointpb.Inventory_SoftwarePackage, len(pkgs.Apt))
		for i, pkg := range pkgs.Apt {
			temp[i] = &agentendpointpb.Inventory_SoftwarePackage{
				Details: formatAptPackage(pkg),
			}
		}
		softwarePackages = append(softwarePackages, temp...)
	}
	if pkgs.Deb != nil {
		temp := make([]*agentendpointpb.Inventory_SoftwarePackage, len(pkgs.Deb))
		for i, pkg := range pkgs.Deb {
			temp[i] = &agentendpointpb.Inventory_SoftwarePackage{
				Details: formatAptPackage(pkg),
			}
		}
		softwarePackages = append(softwarePackages, temp...)
	}
	if pkgs.GooGet != nil {
		temp := make([]*agentendpointpb.Inventory_SoftwarePackage, len(pkgs.GooGet))
		for i, pkg := range pkgs.GooGet {
			temp[i] = &agentendpointpb.Inventory_SoftwarePackage{
				Details: formatGooGetPackage(pkg),
			}
		}
		softwarePackages = append(softwarePackages, temp...)
	}
	if pkgs.Yum != nil {
		temp := make([]*agentendpointpb.Inventory_SoftwarePackage, len(pkgs.Yum))
		for i, pkg := range pkgs.Yum {
			temp[i] = &agentendpointpb.Inventory_SoftwarePackage{
				Details: formatYumPackage(pkg),
			}
		}
		softwarePackages = append(softwarePackages, temp...)
	}
	if pkgs.Zypper != nil {
		temp := make([]*agentendpointpb.Inventory_SoftwarePackage, len(pkgs.Zypper))
		for i, pkg := range pkgs.Zypper {
			temp[i] = &agentendpointpb.Inventory_SoftwarePackage{
				Details: formatZypperPackage(pkg),
			}
		}
		softwarePackages = append(softwarePackages, temp...)
	}
	if pkgs.Rpm != nil {
		temp := make([]*agentendpointpb.Inventory_SoftwarePackage, len(pkgs.Rpm))
		if packages.YumExists {
			for i, pkg := range pkgs.Rpm {
				temp[i] = &agentendpointpb.Inventory_SoftwarePackage{
					Details: formatYumPackage(pkg),
				}
			}
		} else if packages.ZypperExists {
			for i, pkg := range pkgs.Rpm {
				temp[i] = &agentendpointpb.Inventory_SoftwarePackage{
					Details: formatZypperPackage(pkg),
				}
			}
		}
		softwarePackages = append(softwarePackages, temp...)
	}
	if pkgs.ZypperPatches != nil {
		temp := make([]*agentendpointpb.Inventory_SoftwarePackage, len(pkgs.ZypperPatches))
		for i, pkg := range pkgs.ZypperPatches {
			temp[i] = &agentendpointpb.Inventory_SoftwarePackage{
				Details: formatZypperPatch(pkg),
			}
		}
		softwarePackages = append(softwarePackages, temp...)
	}
	if pkgs.WUA != nil {
		temp := make([]*agentendpointpb.Inventory_SoftwarePackage, len(pkgs.WUA))
		for i, pkg := range pkgs.WUA {
			temp[i] = &agentendpointpb.Inventory_SoftwarePackage{
				Details: formatWUAPackage(pkg),
			}
		}
		softwarePackages = append(softwarePackages, temp...)
	}
	if pkgs.QFE != nil {
		temp := make([]*agentendpointpb.Inventory_SoftwarePackage, len(pkgs.QFE))
		for i, pkg := range pkgs.QFE {
			temp[i] = &agentendpointpb.Inventory_SoftwarePackage{
				Details: formatQFEPackage(ctx, pkg),
			}
		}
		softwarePackages = append(softwarePackages, temp...)
	}
	if pkgs.COS != nil {
		temp := make([]*agentendpointpb.Inventory_SoftwarePackage, len(pkgs.COS))
		for i, pkg := range pkgs.COS {
			temp[i] = &agentendpointpb.Inventory_SoftwarePackage{
				Details: formatCOSPackage(pkg),
			}
		}
		softwarePackages = append(softwarePackages, temp...)
	}
	if pkgs.WindowsApplication != nil {
		temp := make([]*agentendpointpb.Inventory_SoftwarePackage, len(pkgs.WindowsApplication))
		for i, pkg := range pkgs.WindowsApplication {
			temp[i] = &agentendpointpb.Inventory_SoftwarePackage{
				Details: formatWindowsApplication(pkg),
			}
		}
		softwarePackages = append(softwarePackages, temp...)
	}
	// Ignore Pip and Gem packages.

	return softwarePackages
}

func formatAptPackage(pkg *packages.PkgInfo) *agentendpointpb.Inventory_SoftwarePackage_AptPackage {
	return &agentendpointpb.Inventory_SoftwarePackage_AptPackage{
		AptPackage: &agentendpointpb.Inventory_VersionedPackage{
			PackageName:  pkg.Name,
			Architecture: pkg.Arch,
			Version:      pkg.Version,
		}}
}

func formatCOSPackage(pkg *packages.PkgInfo) *agentendpointpb.Inventory_SoftwarePackage_CosPackage {
	return &agentendpointpb.Inventory_SoftwarePackage_CosPackage{
		CosPackage: &agentendpointpb.Inventory_VersionedPackage{
			PackageName:  pkg.Name,
			Architecture: pkg.Arch,
			Version:      pkg.Version,
		}}
}

func formatGooGetPackage(pkg *packages.PkgInfo) *agentendpointpb.Inventory_SoftwarePackage_GoogetPackage {
	return &agentendpointpb.Inventory_SoftwarePackage_GoogetPackage{
		GoogetPackage: &agentendpointpb.Inventory_VersionedPackage{
			PackageName:  pkg.Name,
			Architecture: pkg.Arch,
			Version:      pkg.Version,
		}}
}

func formatYumPackage(pkg *packages.PkgInfo) *agentendpointpb.Inventory_SoftwarePackage_YumPackage {
	return &agentendpointpb.Inventory_SoftwarePackage_YumPackage{
		YumPackage: &agentendpointpb.Inventory_VersionedPackage{
			PackageName:  pkg.Name,
			Architecture: pkg.Arch,
			Version:      pkg.Version}}
}

func formatZypperPackage(pkg *packages.PkgInfo) *agentendpointpb.Inventory_SoftwarePackage_ZypperPackage {
	return &agentendpointpb.Inventory_SoftwarePackage_ZypperPackage{
		ZypperPackage: &agentendpointpb.Inventory_VersionedPackage{
			PackageName:  pkg.Name,
			Architecture: pkg.Arch,
			Version:      pkg.Version}}
}

func formatZypperPatch(pkg *packages.ZypperPatch) *agentendpointpb.Inventory_SoftwarePackage_ZypperPatch {
	return &agentendpointpb.Inventory_SoftwarePackage_ZypperPatch{
		ZypperPatch: &agentendpointpb.Inventory_ZypperPatch{
			PatchName: pkg.Name,
			Category:  pkg.Category,
			Severity:  pkg.Severity,
			Summary:   pkg.Summary,
		}}
}

func formatWUAPackage(pkg *packages.WUAPackage) *agentendpointpb.Inventory_SoftwarePackage_WuaPackage {
	var categories []*agentendpointpb.Inventory_WindowsUpdatePackage_WindowsUpdateCategory
	for idx, category := range pkg.Categories {
		categories = append(categories, &agentendpointpb.Inventory_WindowsUpdatePackage_WindowsUpdateCategory{
			Id:   pkg.CategoryIDs[idx],
			Name: category,
		})
	}

	return &agentendpointpb.Inventory_SoftwarePackage_WuaPackage{
		WuaPackage: &agentendpointpb.Inventory_WindowsUpdatePackage{
			Title:                    pkg.Title,
			Description:              pkg.Description,
			Categories:               categories,
			KbArticleIds:             pkg.KBArticleIDs,
			SupportUrl:               pkg.SupportURL,
			MoreInfoUrls:             pkg.MoreInfoURLs,
			UpdateId:                 pkg.UpdateID,
			RevisionNumber:           pkg.RevisionNumber,
			LastDeploymentChangeTime: timestamppb.New(pkg.LastDeploymentChangeTime),
		}}
}

func formatQFEPackage(ctx context.Context, pkg *packages.QFEPackage) *agentendpointpb.Inventory_SoftwarePackage_QfePackage {
	installedTime, err := time.Parse("1/2/2006", pkg.InstalledOn)
	if err != nil {
		clog.Warningf(ctx, "Error parsing QFE InstalledOn date: %v", err)
	}

	return &agentendpointpb.Inventory_SoftwarePackage_QfePackage{
		QfePackage: &agentendpointpb.Inventory_WindowsQuickFixEngineeringPackage{
			Caption:     pkg.Caption,
			Description: pkg.Description,
			HotFixId:    pkg.HotFixID,
			InstallTime: timestamppb.New(installedTime),
		}}
}

func formatWindowsApplication(pkg *packages.WindowsApplication) *agentendpointpb.Inventory_SoftwarePackage_WindowsApplication {

	d := datepb.Date{}
	// We have to check if date is zero.
	// Because zero value of time has Year, Month, Day equal to 1
	if !pkg.InstallDate.IsZero() {
		d = datepb.Date{
			Year:  int32(pkg.InstallDate.Year()),
			Month: int32(pkg.InstallDate.Month()),
			Day:   int32(pkg.InstallDate.Day()),
		}
	}
	return &agentendpointpb.Inventory_SoftwarePackage_WindowsApplication{
		WindowsApplication: &agentendpointpb.Inventory_WindowsApplication{
			DisplayName:    pkg.DisplayName,
			DisplayVersion: pkg.DisplayVersion,
			Publisher:      pkg.Publisher,
			InstallDate:    &d,
			HelpLink:       pkg.HelpLink,
		}}
}
<commit_msg>Don't print raw pointer data. (#446)<commit_after>package agentendpoint

import (
	"context"
	"fmt"
	"reflect"
	"strings"
	"time"

	"github.com/GoogleCloudPlatform/osconfig/agentconfig"
	"github.com/GoogleCloudPlatform/osconfig/attributes"
	"github.com/GoogleCloudPlatform/osconfig/clog"
	"github.com/GoogleCloudPlatform/osconfig/inventory"
	"github.com/GoogleCloudPlatform/osconfig/packages"
	"github.com/GoogleCloudPlatform/osconfig/retryutil"
	"google.golang.org/protobuf/types/known/timestamppb"

	agentendpointpb "google.golang.org/genproto/googleapis/cloud/osconfig/agentendpoint/v1"
	datepb "google.golang.org/genproto/googleapis/type/date"
)

const (
	inventoryURL = agentconfig.ReportURL + "/guestInventory"
)

// ReportInventory writes inventory to guest attributes and reports it to agent endpoint.
func (c *Client) ReportInventory(ctx context.Context) {
	state := inventory.Get(ctx)
	if !agentconfig.DisableInventoryWrite() {
		write(ctx, state, inventoryURL)
	}
	c.report(ctx, state)
}

func write(ctx context.Context, state *inventory.InstanceInventory, url string) {
	clog.Debugf(ctx, "Writing instance inventory to guest attributes.")

	e := reflect.ValueOf(state).Elem()
	t := e.Type()
	for i := 0; i < e.NumField(); i++ {
		f := e.Field(i)
		u := fmt.Sprintf("%s/%s", url, t.Field(i).Name)
		switch f.Kind() {
		case reflect.String:
			clog.Debugf(ctx, "postAttribute %s: %+v", u, f)
			if err := attributes.PostAttribute(u, strings.NewReader(f.String())); err != nil {
				clog.Errorf(ctx, "postAttribute error: %v", err)
			}
		case reflect.Ptr:
			switch reflect.Indirect(f).Kind() {
			case reflect.Struct:
				clog.Debugf(ctx, "postAttributeCompressed %s", u)
				if err := attributes.PostAttributeCompressed(u, f.Interface()); err != nil {
					clog.Errorf(ctx, "postAttributeCompressed error: %v", err)
				}
			}
		}
	}
}

func (c *Client) report(ctx context.Context, state *inventory.InstanceInventory) {
	clog.Debugf(ctx, "Reporting instance inventory to agent endpoint.")
	inventory := formatInventory(ctx, state)

	reportFull := false
	var res *agentendpointpb.ReportInventoryResponse
	var err error
	f := func() error {
		res, err = c.reportInventory(ctx, inventory, reportFull)
		if err != nil {
			return err
		}
		return nil
	}

	if err = retryutil.RetryAPICall(ctx, apiRetrySec*time.Second, "ReportInventory", f); err != nil {
		clog.Errorf(ctx, "Error reporting inventory checksum: %v", err)
		return
	}

	if res.GetReportFullInventory() {
		reportFull = true
		if err = retryutil.RetryAPICall(ctx, apiRetrySec*time.Second, "ReportInventory", f); err != nil {
			clog.Errorf(ctx, "Error reporting full inventory: %v", err)
			return
		}
	}
}

func formatInventory(ctx context.Context, state *inventory.InstanceInventory) *agentendpointpb.Inventory {
	osInfo := &agentendpointpb.Inventory_OsInfo{
		Hostname:             state.Hostname,
		LongName:             state.LongName,
		ShortName:            state.ShortName,
		Version:              state.Version,
		Architecture:         state.Architecture,
		KernelVersion:        state.KernelVersion,
		KernelRelease:        state.KernelRelease,
		OsconfigAgentVersion: state.OSConfigAgentVersion,
	}
	installedPackages := formatPackages(ctx, state.InstalledPackages, state.ShortName)
	availablePackages := formatPackages(ctx, state.PackageUpdates, state.ShortName)

	return &agentendpointpb.Inventory{OsInfo: osInfo, InstalledPackages: installedPackages, AvailablePackages: availablePackages}
}

func formatPackages(ctx context.Context, pkgs *packages.Packages, shortName string) []*agentendpointpb.Inventory_SoftwarePackage {
	var softwarePackages []*agentendpointpb.Inventory_SoftwarePackage
	if pkgs == nil {
		return softwarePackages
	}
	if pkgs.Apt != nil {
		temp := make([]*agentendpointpb.Inventory_SoftwarePackage, len(pkgs.Apt))
		for i, pkg := range pkgs.Apt {
			temp[i] = &agentendpointpb.Inventory_SoftwarePackage{
				Details: formatAptPackage(pkg),
			}
		}
		softwarePackages = append(softwarePackages, temp...)
	}
	if pkgs.Deb != nil {
		temp := make([]*agentendpointpb.Inventory_SoftwarePackage, len(pkgs.Deb))
		for i, pkg := range pkgs.Deb {
			temp[i] = &agentendpointpb.Inventory_SoftwarePackage{
				Details: formatAptPackage(pkg),
			}
		}
		softwarePackages = append(softwarePackages, temp...)
	}
	if pkgs.GooGet != nil {
		temp := make([]*agentendpointpb.Inventory_SoftwarePackage, len(pkgs.GooGet))
		for i, pkg := range pkgs.GooGet {
			temp[i] = &agentendpointpb.Inventory_SoftwarePackage{
				Details: formatGooGetPackage(pkg),
			}
		}
		softwarePackages = append(softwarePackages, temp...)
	}
	if pkgs.Yum != nil {
		temp := make([]*agentendpointpb.Inventory_SoftwarePackage, len(pkgs.Yum))
		for i, pkg := range pkgs.Yum {
			temp[i] = &agentendpointpb.Inventory_SoftwarePackage{
				Details: formatYumPackage(pkg),
			}
		}
		softwarePackages = append(softwarePackages, temp...)
	}
	if pkgs.Zypper != nil {
		temp := make([]*agentendpointpb.Inventory_SoftwarePackage, len(pkgs.Zypper))
		for i, pkg := range pkgs.Zypper {
			temp[i] = &agentendpointpb.Inventory_SoftwarePackage{
				Details: formatZypperPackage(pkg),
			}
		}
		softwarePackages = append(softwarePackages, temp...)
	}
	if pkgs.Rpm != nil {
		temp := make([]*agentendpointpb.Inventory_SoftwarePackage, len(pkgs.Rpm))
		if packages.YumExists {
			for i, pkg := range pkgs.Rpm {
				temp[i] = &agentendpointpb.Inventory_SoftwarePackage{
					Details: formatYumPackage(pkg),
				}
			}
		} else if packages.ZypperExists {
			for i, pkg := range pkgs.Rpm {
				temp[i] = &agentendpointpb.Inventory_SoftwarePackage{
					Details: formatZypperPackage(pkg),
				}
			}
		}
		softwarePackages = append(softwarePackages, temp...)
	}
	if pkgs.ZypperPatches != nil {
		temp := make([]*agentendpointpb.Inventory_SoftwarePackage, len(pkgs.ZypperPatches))
		for i, pkg := range pkgs.ZypperPatches {
			temp[i] = &agentendpointpb.Inventory_SoftwarePackage{
				Details: formatZypperPatch(pkg),
			}
		}
		softwarePackages = append(softwarePackages, temp...)
	}
	if pkgs.WUA != nil {
		temp := make([]*agentendpointpb.Inventory_SoftwarePackage, len(pkgs.WUA))
		for i, pkg := range pkgs.WUA {
			temp[i] = &agentendpointpb.Inventory_SoftwarePackage{
				Details: formatWUAPackage(pkg),
			}
		}
		softwarePackages = append(softwarePackages, temp...)
	}
	if pkgs.QFE != nil {
		temp := make([]*agentendpointpb.Inventory_SoftwarePackage, len(pkgs.QFE))
		for i, pkg := range pkgs.QFE {
			temp[i] = &agentendpointpb.Inventory_SoftwarePackage{
				Details: formatQFEPackage(ctx, pkg),
			}
		}
		softwarePackages = append(softwarePackages, temp...)
	}
	if pkgs.COS != nil {
		temp := make([]*agentendpointpb.Inventory_SoftwarePackage, len(pkgs.COS))
		for i, pkg := range pkgs.COS {
			temp[i] = &agentendpointpb.Inventory_SoftwarePackage{
				Details: formatCOSPackage(pkg),
			}
		}
		softwarePackages = append(softwarePackages, temp...)
	}
	if pkgs.WindowsApplication != nil {
		temp := make([]*agentendpointpb.Inventory_SoftwarePackage, len(pkgs.WindowsApplication))
		for i, pkg := range pkgs.WindowsApplication {
			temp[i] = &agentendpointpb.Inventory_SoftwarePackage{
				Details: formatWindowsApplication(pkg),
			}
		}
		softwarePackages = append(softwarePackages, temp...)
	}
	// Ignore Pip and Gem packages.

	return softwarePackages
}

func formatAptPackage(pkg *packages.PkgInfo) *agentendpointpb.Inventory_SoftwarePackage_AptPackage {
	return &agentendpointpb.Inventory_SoftwarePackage_AptPackage{
		AptPackage: &agentendpointpb.Inventory_VersionedPackage{
			PackageName:  pkg.Name,
			Architecture: pkg.Arch,
			Version:      pkg.Version,
		}}
}

func formatCOSPackage(pkg *packages.PkgInfo) *agentendpointpb.Inventory_SoftwarePackage_CosPackage {
	return &agentendpointpb.Inventory_SoftwarePackage_CosPackage{
		CosPackage: &agentendpointpb.Inventory_VersionedPackage{
			PackageName:  pkg.Name,
			Architecture: pkg.Arch,
			Version:      pkg.Version,
		}}
}

func formatGooGetPackage(pkg *packages.PkgInfo) *agentendpointpb.Inventory_SoftwarePackage_GoogetPackage {
	return &agentendpointpb.Inventory_SoftwarePackage_GoogetPackage{
		GoogetPackage: &agentendpointpb.Inventory_VersionedPackage{
			PackageName:  pkg.Name,
			Architecture: pkg.Arch,
			Version:      pkg.Version,
		}}
}

func formatYumPackage(pkg *packages.PkgInfo) *agentendpointpb.Inventory_SoftwarePackage_YumPackage {
	return &agentendpointpb.Inventory_SoftwarePackage_YumPackage{
		YumPackage: &agentendpointpb.Inventory_VersionedPackage{
			PackageName:  pkg.Name,
			Architecture: pkg.Arch,
			Version:      pkg.Version}}
}

func formatZypperPackage(pkg *packages.PkgInfo) *agentendpointpb.Inventory_SoftwarePackage_ZypperPackage {
	return &agentendpointpb.Inventory_SoftwarePackage_ZypperPackage{
		ZypperPackage: &agentendpointpb.Inventory_VersionedPackage{
			PackageName:  pkg.Name,
			Architecture: pkg.Arch,
			Version:      pkg.Version}}
}

func formatZypperPatch(pkg *packages.ZypperPatch) *agentendpointpb.Inventory_SoftwarePackage_ZypperPatch {
	return &agentendpointpb.Inventory_SoftwarePackage_ZypperPatch{
		ZypperPatch: &agentendpointpb.Inventory_ZypperPatch{
			PatchName: pkg.Name,
			Category:  pkg.Category,
			Severity:  pkg.Severity,
			Summary:   pkg.Summary,
		}}
}

func formatWUAPackage(pkg *packages.WUAPackage) *agentendpointpb.Inventory_SoftwarePackage_WuaPackage {
	var categories []*agentendpointpb.Inventory_WindowsUpdatePackage_WindowsUpdateCategory
	for idx, category := range pkg.Categories {
		categories = append(categories, &agentendpointpb.Inventory_WindowsUpdatePackage_WindowsUpdateCategory{
			Id:   pkg.CategoryIDs[idx],
			Name: category,
		})
	}

	return &agentendpointpb.Inventory_SoftwarePackage_WuaPackage{
		WuaPackage: &agentendpointpb.Inventory_WindowsUpdatePackage{
			Title:                    pkg.Title,
			Description:              pkg.Description,
			Categories:               categories,
			KbArticleIds:             pkg.KBArticleIDs,
			SupportUrl:               pkg.SupportURL,
			MoreInfoUrls:             pkg.MoreInfoURLs,
			UpdateId:                 pkg.UpdateID,
			RevisionNumber:           pkg.RevisionNumber,
			LastDeploymentChangeTime: timestamppb.New(pkg.LastDeploymentChangeTime),
		}}
}

func formatQFEPackage(ctx context.Context, pkg *packages.QFEPackage) *agentendpointpb.Inventory_SoftwarePackage_QfePackage {
	installedTime, err := time.Parse("1/2/2006", pkg.InstalledOn)
	if err != nil {
		clog.Warningf(ctx, "Error parsing QFE InstalledOn date: %v", err)
	}

	return &agentendpointpb.Inventory_SoftwarePackage_QfePackage{
		QfePackage: &agentendpointpb.Inventory_WindowsQuickFixEngineeringPackage{
			Caption:     pkg.Caption,
			Description: pkg.Description,
			HotFixId:    pkg.HotFixID,
			InstallTime: timestamppb.New(installedTime),
		}}
}

func formatWindowsApplication(pkg *packages.WindowsApplication) *agentendpointpb.Inventory_SoftwarePackage_WindowsApplication {

	d := datepb.Date{}
	// We have to check if date is zero.
	// Because zero value of time has Year, Month, Day equal to 1
	if !pkg.InstallDate.IsZero() {
		d = datepb.Date{
			Year:  int32(pkg.InstallDate.Year()),
			Month: int32(pkg.InstallDate.Month()),
			Day:   int32(pkg.InstallDate.Day()),
		}
	}
	return &agentendpointpb.Inventory_SoftwarePackage_WindowsApplication{
		WindowsApplication: &agentendpointpb.Inventory_WindowsApplication{
			DisplayName:    pkg.DisplayName,
			DisplayVersion: pkg.DisplayVersion,
			Publisher:      pkg.Publisher,
			InstallDate:    &d,
			HelpLink:       pkg.HelpLink,
		}}
}
<|endoftext|>


<commit_before>package ircclient

// Handles basic IRC protocol messages (like PING)

import (
	"log"
)

type basicProtocol struct {
	ic *IRCClient
}

func (bp *basicProtocol) Register(cl *IRCClient) {
	bp.ic = cl
}

func (bp *basicProtocol) String() string {
	return "basic"
}

func (bp *basicProtocol) Usage(cmd string) string {
	// stub, no commands here
	return ""
}

func (bp *basicProtocol) ProcessLine(msg *IRCMessage) {
	switch msg.Command {
	case "PING":
		if len(msg.Args) != 1 {
			log.Printf("WARNING: Invalid PING received")
		}
		bp.ic.conn.Output <- "PONG :" + msg.Args[0]
	}
}
func (bp *basicProtocol) Unregister() {
}

func (bp *basicProtocol) Info() string {
	return "basic irc protocol (e.g. PING), implemented as plugin."
}

func (bp *basicProtocol) ProcessCommand(cmd *IRCCommand) {
}
<commit_msg>add a ping routine<commit_after>package ircclient

// Handles basic IRC protocol messages (like PING)

import (
	"log"
	"time"
)

type basicProtocol struct {
	ic          *IRCClient
	chanTimeout chan *IRCMessage
}

func (bp *basicProtocol) Register(cl *IRCClient) {
	bp.ic = cl
	bp.chanTimeout = make(chan *IRCMessage)

	// start pinging routine
	go func() {
		for {
			time.Sleep(2 * time.Minute)
			bp.ic.SendLine("PING :" + bp.ic.GetStringOption("Server", "nick"))

			select {
			case <-bp.chanTimeout:
			case <-time.After(20 * time.Second):
				log.Println("Ping timeout")
				bp.ic.Disconnect("Ping timeout")
			}

		}
	}()
}

func (bp *basicProtocol) String() string {
	return "basic"
}

func (bp *basicProtocol) Usage(cmd string) string {
	// stub, no commands here
	return ""
}

func (bp *basicProtocol) ProcessLine(msg *IRCMessage) {
	switch msg.Command {
	case "PING":
		if len(msg.Args) != 1 {
			log.Printf("WARNING: Invalid PING received")
		}
		bp.ic.SendLine("PONG :" + msg.Args[0])
	case "PONG":
		bp.chanTimeout <- msg
	}
}
func (bp *basicProtocol) Unregister() {
}

func (bp *basicProtocol) Info() string {
	return "basic irc protocol (e.g. PING), implemented as plugin."
}

func (bp *basicProtocol) ProcessCommand(cmd *IRCCommand) {
}
<|endoftext|>


<commit_before>package main

import (
	"code.google.com/p/gcfg"
	"database/sql"
	"flag"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"github.com/streadway/amqp"
	"io/ioutil"
	"log"
	"net/http"
	"strconv"
	"time"
)

var (
	debug = flag.Bool("d", false, "turn on debug info")
	cfg   struct {
		DB struct {
			Driver string
			Dsn    string
			Table  string
		}
		RabbitMQ struct {
			Uri      string
			Exchange string
		}
		HTTP struct {
			Host string
			Port string
		}
	}

	db     *sql.DB
	broker *amqp.Connection
	wakeUp = make(chan int, 1)
)

type Job struct {
	routingKey string
	body       string
	interval   uint
	nextRun    time.Time
	state      string
}

func handleSchedule(w http.ResponseWriter, r *http.Request) {
	routingKey, body, interval_s :=
		r.FormValue("routing_key"), r.FormValue("body"), r.FormValue("interval")
	log.Println("/schedule", routingKey, body)

	interval, err := strconv.ParseInt(interval_s, 10, 64)
	if err != nil {
		panic(err)
	}

	next_run := time.Now().UTC().Add(time.Duration(interval) * time.Second)
	_, err = db.Exec("INSERT INTO "+cfg.DB.Table+" "+
		"(routing_key, body, `interval`, next_run, state) "+
		"VALUES(?, ?, ?, ?, 'WAITING') "+
		"ON DUPLICATE KEY UPDATE `interval`=?",
		routingKey, body, interval, next_run, interval)
	if err != nil {
		panic(err)
	}

	// Wake up the publisher
	//
	// publisher() may be sleeping for the next job on the queue
	// at the time we schedule a new Job. Let it wake up so it can
	// re-fetch the new Job from the front of the queue.
	//
	// The code below is an idiom for non-blocking send to a channel
	select {
	case wakeUp <- 1:
		log.Println("Sent wakeup signal")
	default:
		log.Println("Skipped wakeup signal")
	}
}

func handleCancel(w http.ResponseWriter, r *http.Request) {
	routingKey, body := r.FormValue("routing_key"), r.FormValue("body")
	log.Println("/cancel", routingKey, body)

	_, err := db.Exec("DELETE FROM "+cfg.DB.Table+" "+
		"WHERE routing_key=? AND body=?", routingKey, body)
	if err != nil {
		panic(err)
	}
}

func front() (*Job, error) {
	j := Job{}
	row := db.QueryRow("SELECT routing_key, body, `interval`, next_run, state " +
		"FROM " + cfg.DB.Table + " " +
		"WHERE next_run = (" +
		"SELECT MIN(next_run) FROM " + cfg.DB.Table + ")")
	err := row.Scan(&j.routingKey, &j.body, &j.interval, &j.nextRun, &j.state)
	if err != nil {
		return nil, err
	}
	return &j, nil
}

func (j *Job) Delete() error {
	_, err := db.Exec("DELETE FROM "+cfg.DB.Table+" "+
		"WHERE routing_key=? AND body=?",
		j.routingKey, j.body)
	if err != nil {
		return err
	}
	return nil
}

func (j *Job) Publish() error {
	log.Println("publish", *j)
	return nil
}

func (j *Job) Remaining() time.Duration {
	return -time.Since(j.nextRun)
}

func publisher() {
	for {
		job, err := front()
		if err != nil {
			fmt.Println(err)
			log.Println("Waiting wakeup signal")
			<-wakeUp
			log.Println("Got wakeup signal")
			continue
		}
		log.Println("Next job:", job, "Remaining:", job.Remaining())

		now := time.Now().UTC()
		if job.nextRun.After(now) {
			select {
			case <-time.After(job.Remaining()):
			case <-wakeUp:
				continue
			}
		} else {
			job.Publish()
			err = job.Delete()
			if err != nil {
				log.Println(err)
			}
		}

	}
}

func main() {
	// Setup logging
	flag.Parse()
	if !*debug {
		log.SetOutput(ioutil.Discard)
	}

	// Read config
	err := gcfg.ReadFileInto(&cfg, "dalga.ini")
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("Read config: ", cfg)

	// Connect to database
	db, err = sql.Open(cfg.DB.Driver, cfg.DB.Dsn)
	if err != nil {
		log.Fatal(err)
	}
	err = db.Ping()
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("Connected to DB")

	// Connect to RabbitMQ
	_, err = amqp.Dial(cfg.RabbitMQ.Uri)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("Connected to RabbitMQ")

	// Run publisher
	go publisher()

	// Start HTTP server
	addr := cfg.HTTP.Host + ":" + cfg.HTTP.Port
	http.HandleFunc("/schedule", handleSchedule)
	http.HandleFunc("/cancel", handleCancel)
	http.ListenAndServe(addr, nil)
}
<commit_msg>add comment<commit_after>package main

import (
	"code.google.com/p/gcfg"
	"database/sql"
	"flag"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"github.com/streadway/amqp"
	"io/ioutil"
	"log"
	"net/http"
	"strconv"
	"time"
)

var (
	debug = flag.Bool("d", false, "turn on debug info")
	cfg   struct {
		DB struct {
			Driver string
			Dsn    string
			Table  string
		}
		RabbitMQ struct {
			Uri      string
			Exchange string
		}
		HTTP struct {
			Host string
			Port string
		}
	}

	db     *sql.DB
	broker *amqp.Connection
	wakeUp = make(chan int, 1)
)

type Job struct {
	routingKey string
	body       string
	interval   uint
	nextRun    time.Time
	state      string
}

func handleSchedule(w http.ResponseWriter, r *http.Request) {
	routingKey, body, interval_s :=
		r.FormValue("routing_key"), r.FormValue("body"), r.FormValue("interval")
	log.Println("/schedule", routingKey, body)

	interval, err := strconv.ParseInt(interval_s, 10, 64)
	if err != nil {
		panic(err)
	}

	next_run := time.Now().UTC().Add(time.Duration(interval) * time.Second)
	_, err = db.Exec("INSERT INTO "+cfg.DB.Table+" "+
		"(routing_key, body, `interval`, next_run, state) "+
		"VALUES(?, ?, ?, ?, 'WAITING') "+
		"ON DUPLICATE KEY UPDATE `interval`=?",
		routingKey, body, interval, next_run, interval)
	if err != nil {
		panic(err)
	}

	// Wake up the publisher
	//
	// publisher() may be sleeping for the next job on the queue
	// at the time we schedule a new Job. Let it wake up so it can
	// re-fetch the new Job from the front of the queue.
	//
	// The code below is an idiom for non-blocking send to a channel
	select {
	case wakeUp <- 1:
		log.Println("Sent wakeup signal")
	default:
		log.Println("Skipped wakeup signal")
	}
}

func handleCancel(w http.ResponseWriter, r *http.Request) {
	routingKey, body := r.FormValue("routing_key"), r.FormValue("body")
	log.Println("/cancel", routingKey, body)

	_, err := db.Exec("DELETE FROM "+cfg.DB.Table+" "+
		"WHERE routing_key=? AND body=?", routingKey, body)
	if err != nil {
		panic(err)
	}
}

func front() (*Job, error) {
	j := Job{}
	row := db.QueryRow("SELECT routing_key, body, `interval`, next_run, state " +
		"FROM " + cfg.DB.Table + " " +
		"WHERE next_run = (" +
		"SELECT MIN(next_run) FROM " + cfg.DB.Table + ")")
	err := row.Scan(&j.routingKey, &j.body, &j.interval, &j.nextRun, &j.state)
	if err != nil {
		return nil, err
	}
	return &j, nil
}

func (j *Job) Delete() error {
	_, err := db.Exec("DELETE FROM "+cfg.DB.Table+" "+
		"WHERE routing_key=? AND body=?",
		j.routingKey, j.body)
	if err != nil {
		return err
	}
	return nil
}

func (j *Job) Publish() error {
	log.Println("publish", *j)
	return nil
}

func (j *Job) Remaining() time.Duration {
	return -time.Since(j.nextRun)
}

func publisher() {
	for {
		job, err := front()
		if err != nil {
			fmt.Println(err)
			log.Println("Waiting wakeup signal")
			<-wakeUp
			log.Println("Got wakeup signal")
			continue
		}
		log.Println("Next job:", job, "Remaining:", job.Remaining())

		now := time.Now().UTC()
		if job.nextRun.After(now) {
			// Wait until the next Job time or
			// the webserver's /schedule handler wakes us up
			select {
			case <-time.After(job.Remaining()):
			case <-wakeUp:
				continue
			}
		} else {
			job.Publish()
			err = job.Delete()
			if err != nil {
				log.Println(err)
			}
		}

	}
}

func main() {
	// Setup logging
	flag.Parse()
	if !*debug {
		log.SetOutput(ioutil.Discard)
	}

	// Read config
	err := gcfg.ReadFileInto(&cfg, "dalga.ini")
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("Read config: ", cfg)

	// Connect to database
	db, err = sql.Open(cfg.DB.Driver, cfg.DB.Dsn)
	if err != nil {
		log.Fatal(err)
	}
	err = db.Ping()
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("Connected to DB")

	// Connect to RabbitMQ
	_, err = amqp.Dial(cfg.RabbitMQ.Uri)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("Connected to RabbitMQ")

	// Run publisher
	go publisher()

	// Start HTTP server
	addr := cfg.HTTP.Host + ":" + cfg.HTTP.Port
	http.HandleFunc("/schedule", handleSchedule)
	http.HandleFunc("/cancel", handleCancel)
	http.ListenAndServe(addr, nil)
}
<|endoftext|>


<commit_before>// Copyright 2016 Hajime Hoshi
//
// 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 wav provides WAV (RIFF) decoder.
package wav

import (
	"bytes"
	"fmt"
	"io"

	"github.com/hajimehoshi/ebiten/audio"
)

// Stream is a decoded audio stream.
type Stream struct {
	src        audio.ReadSeekCloser
	headerSize int64
	dataSize   int64
	remaining  int64
}

// Read is implementation of io.Reader's Read.
func (s *Stream) Read(p []byte) (int, error) {
	if s.remaining <= 0 {
		return 0, io.EOF
	}
	if s.remaining < int64(len(p)) {
		p = p[0:s.remaining]
	}
	n, err := s.src.Read(p)
	s.remaining -= int64(n)
	return n, err
}

// Seek is implementation of io.Seeker's Seek.
func (s *Stream) Seek(offset int64, whence int) (int64, error) {
	if whence == io.SeekStart {
		offset += s.headerSize
	}
	n, err := s.src.Seek(offset, whence)
	if err != nil {
		return 0, err
	}
	if n-s.headerSize < 0 {
		return 0, fmt.Errorf("wav: invalid offset")
	}
	s.remaining = s.dataSize - (n - s.headerSize)
	return n - s.headerSize, nil
}

// Read is implementation of io.Closer's Close.
func (s *Stream) Close() error {
	return s.src.Close()
}

// Size returns the size of decoded stream in bytes.
func (s *Stream) Size() int64 {
	return s.dataSize
}

// Decode decodes WAV (RIFF) data to playable stream.
//
// The format must be 2 channels, 16bit little endian PCM.
// The sample rate must be same as that of audio context.
func Decode(context *audio.Context, src audio.ReadSeekCloser) (*Stream, error) {
	buf := make([]byte, 12)
	n, err := io.ReadFull(src, buf)
	if n != len(buf) {
		return nil, fmt.Errorf("wav: invalid header")
	}
	if err != nil {
		return nil, err
	}
	if !bytes.Equal(buf[0:4], []byte("RIFF")) {
		return nil, fmt.Errorf("wav: invalid header: 'RIFF' not found")
	}
	if !bytes.Equal(buf[8:12], []byte("WAVE")) {
		return nil, fmt.Errorf("wav: invalid header: 'WAVE' not found")
	}

	// Read chunks
	dataSize := int64(0)
	headerSize := int64(0)
chunks:
	for {
		buf := make([]byte, 8)
		n, err := io.ReadFull(src, buf)
		if n != len(buf) {
			return nil, fmt.Errorf("wav: invalid header")
		}
		if err != nil {
			return nil, err
		}
		headerSize += 8
		size := int64(buf[4]) | int64(buf[5])<<8 | int64(buf[6])<<16 | int64(buf[7])<<24
		switch {
		case bytes.Equal(buf[0:4], []byte("fmt ")):
			if size != 16 {
				return nil, fmt.Errorf("wav: invalid header: maybe non-PCM file?")
			}
			buf := make([]byte, size)
			n, err := io.ReadFull(src, buf)
			if n != len(buf) {
				return nil, fmt.Errorf("wav: invalid header")
			}
			if err != nil {
				return nil, err
			}
			format := int(buf[0]) | int(buf[1])<<8
			if format != 1 {
				return nil, fmt.Errorf("wav: format must be linear PCM")
			}
			channelNum := int(buf[2]) | int(buf[3])<<8
			// TODO: Remove this magic number
			if channelNum != 2 {
				return nil, fmt.Errorf("wav: channel num must be 2")
			}
			bitsPerSample := int(buf[14]) | int(buf[15])<<8
			// TODO: Remove this magic number
			if bitsPerSample != 16 {
				return nil, fmt.Errorf("wav: bits per sample must be 16")
			}
			sampleRate := int64(buf[4]) | int64(buf[5])<<8 | int64(buf[6])<<16 | int64(buf[7])<<24
			if int64(context.SampleRate()) != sampleRate {
				return nil, fmt.Errorf("wav: sample rate must be %d but %d", context.SampleRate(), sampleRate)
			}
			headerSize += size
		case bytes.Equal(buf[0:4], []byte("data")):
			dataSize = size
			break chunks
		default:
			buf := make([]byte, size)
			n, err := io.ReadFull(src, buf)
			if n != len(buf) {
				return nil, fmt.Errorf("wav: invalid header")
			}
			if err != nil {
				return nil, err
			}
			headerSize += size
		}
	}
	s := &Stream{
		src:        src,
		headerSize: headerSize,
		dataSize:   dataSize,
		remaining:  dataSize,
	}
	return s, nil
}
<commit_msg>audio: Bug fix: Fix Seek: wav file can have a tail<commit_after>// Copyright 2016 Hajime Hoshi
//
// 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 wav provides WAV (RIFF) decoder.
package wav

import (
	"bytes"
	"fmt"
	"io"

	"github.com/hajimehoshi/ebiten/audio"
)

// Stream is a decoded audio stream.
type Stream struct {
	src        audio.ReadSeekCloser
	headerSize int64
	dataSize   int64
	remaining  int64
}

// Read is implementation of io.Reader's Read.
func (s *Stream) Read(p []byte) (int, error) {
	if s.remaining <= 0 {
		return 0, io.EOF
	}
	if s.remaining < int64(len(p)) {
		p = p[0:s.remaining]
	}
	n, err := s.src.Read(p)
	s.remaining -= int64(n)
	return n, err
}

// Seek is implementation of io.Seeker's Seek.
func (s *Stream) Seek(offset int64, whence int) (int64, error) {
	if whence == io.SeekStart {
		offset += s.headerSize
	}
	n, err := s.src.Seek(offset, whence)
	if err != nil {
		return 0, err
	}
	if n-s.headerSize < 0 {
		return 0, fmt.Errorf("wav: invalid offset")
	}
	s.remaining = s.dataSize - (n - s.headerSize)
	// There could be a tail in wav file.
	if s.remaining < 0 {
		s.remaining = 0
		return s.dataSize, nil
	}
	return n - s.headerSize, nil
}

// Read is implementation of io.Closer's Close.
func (s *Stream) Close() error {
	return s.src.Close()
}

// Size returns the size of decoded stream in bytes.
func (s *Stream) Size() int64 {
	return s.dataSize
}

// Decode decodes WAV (RIFF) data to playable stream.
//
// The format must be 2 channels, 16bit little endian PCM.
// The sample rate must be same as that of audio context.
func Decode(context *audio.Context, src audio.ReadSeekCloser) (*Stream, error) {
	buf := make([]byte, 12)
	n, err := io.ReadFull(src, buf)
	if n != len(buf) {
		return nil, fmt.Errorf("wav: invalid header")
	}
	if err != nil {
		return nil, err
	}
	if !bytes.Equal(buf[0:4], []byte("RIFF")) {
		return nil, fmt.Errorf("wav: invalid header: 'RIFF' not found")
	}
	if !bytes.Equal(buf[8:12], []byte("WAVE")) {
		return nil, fmt.Errorf("wav: invalid header: 'WAVE' not found")
	}

	// Read chunks
	dataSize := int64(0)
	headerSize := int64(0)
chunks:
	for {
		buf := make([]byte, 8)
		n, err := io.ReadFull(src, buf)
		if n != len(buf) {
			return nil, fmt.Errorf("wav: invalid header")
		}
		if err != nil {
			return nil, err
		}
		headerSize += 8
		size := int64(buf[4]) | int64(buf[5])<<8 | int64(buf[6])<<16 | int64(buf[7])<<24
		switch {
		case bytes.Equal(buf[0:4], []byte("fmt ")):
			if size != 16 {
				return nil, fmt.Errorf("wav: invalid header: maybe non-PCM file?")
			}
			buf := make([]byte, size)
			n, err := io.ReadFull(src, buf)
			if n != len(buf) {
				return nil, fmt.Errorf("wav: invalid header")
			}
			if err != nil {
				return nil, err
			}
			format := int(buf[0]) | int(buf[1])<<8
			if format != 1 {
				return nil, fmt.Errorf("wav: format must be linear PCM")
			}
			channelNum := int(buf[2]) | int(buf[3])<<8
			// TODO: Remove this magic number
			if channelNum != 2 {
				return nil, fmt.Errorf("wav: channel num must be 2")
			}
			bitsPerSample := int(buf[14]) | int(buf[15])<<8
			// TODO: Remove this magic number
			if bitsPerSample != 16 {
				return nil, fmt.Errorf("wav: bits per sample must be 16")
			}
			sampleRate := int64(buf[4]) | int64(buf[5])<<8 | int64(buf[6])<<16 | int64(buf[7])<<24
			if int64(context.SampleRate()) != sampleRate {
				return nil, fmt.Errorf("wav: sample rate must be %d but %d", context.SampleRate(), sampleRate)
			}
			headerSize += size
		case bytes.Equal(buf[0:4], []byte("data")):
			dataSize = size
			break chunks
		default:
			buf := make([]byte, size)
			n, err := io.ReadFull(src, buf)
			if n != len(buf) {
				return nil, fmt.Errorf("wav: invalid header")
			}
			if err != nil {
				return nil, err
			}
			headerSize += size
		}
	}
	s := &Stream{
		src:        src,
		headerSize: headerSize,
		dataSize:   dataSize,
		remaining:  dataSize,
	}
	return s, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/proxypoke/i3ipc"
	"os"
	"os/exec"
	"strings"
)

func dfsTree(t *i3ipc.I3Node) (names []string, nodes []*i3ipc.I3Node) {
	if t.Geometry.Width > 0 && t.Geometry.Height > 0 && t.Window > 0 && !t.Focused {
		names = append(names, t.Name)
		nodes = append(nodes, t)
	}
	visitChildren := func(children *[]i3ipc.I3Node) {
		for i, _ := range *children {
			newNames, newNodes := dfsTree(&(*children)[i])
			names = append(names, newNames...)
			nodes = append(nodes, newNodes...)
		}
	}
	visitChildren(&t.FloatingNodes)
	visitChildren(&t.Nodes)
	return
}

func checkError(e error) {
	if e != nil {
		fmt.Fprintln(os.Stderr, e)
		os.Exit(0)
	}
}

func main() {
	ipc, e := i3ipc.GetIPCSocket()
	checkError(e)

	tree, e := ipc.GetTree()
	checkError(e)

	names, nodes := dfsTree(&tree)

	cmd := exec.Command("dmenu", os.Args[1:len(os.Args)]...)
	cmd.Stdin = strings.NewReader(strings.Join(names, "\n"))

	out, e := cmd.Output()
	checkError(e)

	name := strings.TrimRight(string(out), "\n")

	var node *i3ipc.I3Node
	for i, s := range names {
		if s == name {
			node = nodes[i]
			break
		}
	}

	if node != nil {
		msg := fmt.Sprint("[con_id=", node.Id, "] focus")
		ipc.Raw(i3ipc.I3Command, msg)
	}
}
<commit_msg>Ignore dockarea<commit_after>package main

import (
	"fmt"
	"github.com/proxypoke/i3ipc"
	"os"
	"os/exec"
	"strings"
)

func dfsTree(t *i3ipc.I3Node) (names []string, nodes []*i3ipc.I3Node) {
	if t.Layout == "dockarea" {
		return
	}

	if t.Geometry.Width > 0 && t.Geometry.Height > 0 && t.Window > 0 && !t.Focused {
		names = append(names, t.Name)
		nodes = append(nodes, t)
	}
	visitChildren := func(children *[]i3ipc.I3Node) {
		for i, _ := range *children {
			newNames, newNodes := dfsTree(&(*children)[i])
			names = append(names, newNames...)
			nodes = append(nodes, newNodes...)
		}
	}
	visitChildren(&t.FloatingNodes)
	visitChildren(&t.Nodes)
	return
}

func checkError(e error) {
	if e != nil {
		fmt.Fprintln(os.Stderr, e)
		os.Exit(0)
	}
}

func main() {
	ipc, e := i3ipc.GetIPCSocket()
	checkError(e)

	tree, e := ipc.GetTree()
	checkError(e)

	names, nodes := dfsTree(&tree)

	cmd := exec.Command("dmenu", os.Args[1:len(os.Args)]...)
	cmd.Stdin = strings.NewReader(strings.Join(names, "\n"))

	out, e := cmd.Output()
	checkError(e)

	name := strings.TrimRight(string(out), "\n")

	var node *i3ipc.I3Node
	for i, s := range names {
		if s == name {
			node = nodes[i]
			break
		}
	}

	if node != nil {
		msg := fmt.Sprint("[con_id=", node.Id, "] focus")
		ipc.Raw(i3ipc.I3Command, msg)
	}
}
<|endoftext|>


<commit_before>package imageutil

import (
	"image"
	"image/color"
)

// ImageReader implements the same methods as the standard image interface.
type ImageReader interface {
	At(x, y int) color.Color
	Bounds() image.Rectangle
	ColorModel() color.Model
}

// ImageWriter implements a method for setting colors at individual
// coordinates (which all standard image types implement).
type ImageWriter interface {
	Set(x, y int, c color.Color)
}

// ImageReadWriter implements both ImageReader and ImageWriter.
type ImageReadWriter interface {
	ImageReader
	ImageWriter
}

// Copy concurrently copies Color values from a source ImageReader to a
// destination ImageReadWriter.
func Copy(dst ImageReadWriter, src ImageReader) {
	if dst != src {
		QuickRP(
			AllPointsRP(
				func(pt image.Point) {
					dst.Set(pt.X, pt.Y, src.At(pt.X, pt.Y))
				},
			),
		)(dst.Bounds().Union(src.Bounds()))
	}
}

// ConvertToRGBA returns a new copy of an ImageReader as a new RGBA using Copy
// to concurrently set Color values.
func ConvertToRGBA(src ImageReader) *image.RGBA {
	dst := image.NewRGBA(src.Bounds())
	Copy(dst, src)
	return dst
}

// ConvertToRGBA64 returns a new copy of an ImageReader as a new RGBA64 using
// Copy to concurrently set Color values.
func ConvertToRGBA64(src ImageReader) *image.RGBA64 {
	dst := image.NewRGBA64(src.Bounds())
	Copy(dst, src)
	return dst
}

// ConvertToNRGBA returns a new copy of an ImageReader as a new NRGBA using
// Copy to concurrently set Color values.
func ConvertToNRGBA(src ImageReader) *image.NRGBA {
	dst := image.NewNRGBA(src.Bounds())
	Copy(dst, src)
	return dst
}

// ConvertToNRGBA64 returns a new copy of an ImageReader as a new NRGBA64
// using Copy to concurrently set Color values.
func ConvertToNRGBA64(src ImageReader) *image.NRGBA64 {
	dst := image.NewNRGBA64(src.Bounds())
	Copy(dst, src)
	return dst
}

// ConvertToAlpha returns a new copy of an ImageReader as a new Alpha using
// Copy to concurrently set Color values.
func ConvertToAlpha(src ImageReader) *image.Alpha {
	dst := image.NewAlpha(src.Bounds())
	Copy(dst, src)
	return dst
}

// ConvertToAlpha16 returns a new copy of an ImageReader as a new Alpha16
// using Copy to concurrently set Color values.
func ConvertToAlpha16(src ImageReader) *image.Alpha16 {
	dst := image.NewAlpha16(src.Bounds())
	Copy(dst, src)
	return dst
}

// ConvertToGray returns a new copy of an ImageReader as a new Gray using Copy
// to concurrently set Color values.
func ConvertToGray(src ImageReader) *image.Gray {
	dst := image.NewGray(src.Bounds())
	Copy(dst, src)
	return dst
}

// ConvertToGray16 returns a new copy of an ImageReader as a new Gray16 using
// Copy to concurrently set Color values.
func ConvertToGray16(src ImageReader) *image.Gray16 {
	dst := image.NewGray16(src.Bounds())
	Copy(dst, src)
	return dst
}
<commit_msg>check type before copy in convert methods, clarify docs<commit_after>package imageutil

import (
	"image"
	"image/color"
)

// ImageReader implements the same methods as the standard image interface.
type ImageReader interface {
	At(x, y int) color.Color
	Bounds() image.Rectangle
	ColorModel() color.Model
}

// ImageWriter implements a method for setting colors at individual
// coordinates (which all standard image types implement).
type ImageWriter interface {
	Set(x, y int, c color.Color)
}

// ImageReadWriter implements both ImageReader and ImageWriter.
type ImageReadWriter interface {
	ImageReader
	ImageWriter
}

// Copy concurrently copies Color values from a source ImageReader to a
// destination ImageReadWriter.
func Copy(dst ImageReadWriter, src ImageReader) {
	if dst != src {
		QuickRP(
			AllPointsRP(
				func(pt image.Point) {
					dst.Set(pt.X, pt.Y, src.At(pt.X, pt.Y))
				},
			),
		)(dst.Bounds().Union(src.Bounds()))
	}
}

// ConvertToRGBA returns an *image.RGBA instance by asserting the given
// ImageReader has that type or, if it does not, using Copy to concurrently
// set the color.Color values of a new *image.RGBA instance with the same
// bounds.
func ConvertToRGBA(src ImageReader) *image.RGBA {
	if dst, ok := src.(*image.RGBA); ok {
		return dst
	}
	dst := image.NewRGBA(src.Bounds())
	Copy(dst, src)
	return dst
}

// ConvertToRGBA64 returns an *image.RGBA64 instance by asserting the given
// ImageReader has that type or, if it does not, using Copy to concurrently
// set the color.Color values of a new *image.RGBA64 instance with the same
// bounds.
func ConvertToRGBA64(src ImageReader) *image.RGBA64 {
	if dst, ok := src.(*image.RGBA64); ok {
		return dst
	}
	dst := image.NewRGBA64(src.Bounds())
	Copy(dst, src)
	return dst
}

// ConvertToNRGBA returns an *image.NRGBA instance by asserting the given
// ImageReader has that type or, if it does not, using Copy to concurrently
// set the color.Color values of a new *image.NRGBA instance with the same
// bounds.
func ConvertToNRGBA(src ImageReader) *image.NRGBA {
	if dst, ok := src.(*image.NRGBA); ok {
		return dst
	}
	dst := image.NewNRGBA(src.Bounds())
	Copy(dst, src)
	return dst
}

// ConvertToNRGBA64 returns an *image.NRGBA64 instance by asserting the given
// ImageReader has that type or, if it does not, using Copy to concurrently
// set the color.Color values of a new *image.NRGBA64 instance with the same
// bounds.
func ConvertToNRGBA64(src ImageReader) *image.NRGBA64 {
	if dst, ok := src.(*image.NRGBA64); ok {
		return dst
	}
	dst := image.NewNRGBA64(src.Bounds())
	Copy(dst, src)
	return dst
}

// ConvertToAlpha returns an *image.Alpha instance by asserting the given
// ImageReader has that type or, if it does not, using Copy to concurrently
// set the color.Color values of a new *image.Alpha instance with the same
// bounds.
func ConvertToAlpha(src ImageReader) *image.Alpha {
	if dst, ok := src.(*image.Alpha); ok {
		return dst
	}
	dst := image.NewAlpha(src.Bounds())
	Copy(dst, src)
	return dst
}

// ConvertToAlpha16 returns an *image.Alpha16 instance by asserting the given
// ImageReader has that type or, if it does not, using Copy to concurrently
// set the color.Color values of a new *image.Alpha16 instance with the same
// bounds.
func ConvertToAlpha16(src ImageReader) *image.Alpha16 {
	if dst, ok := src.(*image.Alpha16); ok {
		return dst
	}
	dst := image.NewAlpha16(src.Bounds())
	Copy(dst, src)
	return dst
}

// ConvertToGray returns an *image.Gray instance by asserting the given
// ImageReader has that type or, if it does not, using Copy to concurrently
// set the color.Color values of a new *image.Gray instance with the same
// bounds.
func ConvertToGray(src ImageReader) *image.Gray {
	if dst, ok := src.(*image.Gray); ok {
		return dst
	}
	dst := image.NewGray(src.Bounds())
	Copy(dst, src)
	return dst
}

// ConvertToGray16 returns an *image.Gray16 instance by asserting the given
// ImageReader has that type or, if it does not, using Copy to concurrently
// set the color.Color values of a new *image.Gray16 instance with the same
// bounds.
func ConvertToGray16(src ImageReader) *image.Gray16 {
	if dst, ok := src.(*image.Gray16); ok {
		return dst
	}
	dst := image.NewGray16(src.Bounds())
	Copy(dst, src)
	return dst
}
<|endoftext|>


<commit_before>package history

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"path"
	"sync"

	"gopkg.in/mgo.v2/bson"

	"github.com/henrylee2cn/pholcus/app/downloader/context"
	"github.com/henrylee2cn/pholcus/common/mgo"
	"github.com/henrylee2cn/pholcus/common/mysql"
	"github.com/henrylee2cn/pholcus/config"
	"github.com/henrylee2cn/pholcus/logs"
)

type Historier interface {
	// 读取成功记录
	ReadSuccess(provider string, inherit bool)
	// 更新或加入成功记录
	UpsertSuccess(Record) bool
	// 删除成功记录
	DeleteSuccess(Record)
	// I/O输出成功记录，但不清缓存
	FlushSuccess(provider string)

	// 读取失败记录
	ReadFailure(provider string, inherit bool)
	// 更新或加入失败记录
	UpsertFailure(*context.Request) bool
	// 删除失败记录
	DeleteFailure(*context.Request)
	// I/O输出失败记录，但不清缓存
	FlushFailure(provider string)
	// 获取指定蜘蛛在上一次运行时失败的请求
	PullFailure(spiderName string) []*context.Request

	// 清空缓存，但不输出
	Empty()
}

type Record interface {
	GetUrl() string
	GetMethod() string
}

var (
	MGO_DB = config.MGO.DB

	SUCCESS_FILE = config.HISTORY.FILE_NAME + "_y"
	FAILURE_FILE = config.HISTORY.FILE_NAME + "_n"

	SUCCESS_FILE_FULL = config.COMM_PATH.CACHE + "/" + SUCCESS_FILE
	FAILURE_FILE_FULL = config.COMM_PATH.CACHE + "/" + FAILURE_FILE
)

type History struct {
	*Success
	*Failure
	provider string
	sync.RWMutex
}

func New() Historier {
	return &History{
		Success: &Success{
			new: make(map[string]bool),
			old: make(map[string]bool),
		},
		Failure: &Failure{
			new: make(map[string]map[string]bool),
			old: make(map[string]map[string]bool),
		},
	}
}

// 读取成功记录
func (self *History) ReadSuccess(provider string, inherit bool) {
	self.RWMutex.Lock()
	self.provider = provider
	self.RWMutex.Unlock()

	if !inherit {
		// 不继承历史记录时
		self.Success.old = make(map[string]bool)
		self.Success.new = make(map[string]bool)
		self.Success.inheritable = false
		return

	} else if self.Success.inheritable {
		// 本次与上次均继承历史记录时
		return

	} else {
		// 上次没有继承历史记录，但本次继承时
		self.Success.old = make(map[string]bool)
		self.Success.new = make(map[string]bool)
		self.Success.inheritable = true
	}

	switch provider {
	case "mgo":
		var docs = map[string]interface{}{}
		err := mgo.Mgo(&docs, "find", map[string]interface{}{
			"Database":   MGO_DB,
			"Collection": SUCCESS_FILE,
		})
		if err != nil {
			logs.Log.Error("从mgo读取成功记录: %v", err)
			return
		}
		for _, v := range docs["Docs"].([]interface{}) {
			self.Success.old[v.(bson.M)["_id"].(string)] = true
		}

	case "mysql":
		db, ok := mysql.MysqlPool.GetOne().(*mysql.MysqlSrc)
		if !ok || db == nil {
			// logs.Log.Error("链接Mysql数据库超时，无法读取成功记录！")
			return
		}
		defer mysql.MysqlPool.Free(db)
		rows, err := mysql.New(db.DB).
			SetTableName("`" + SUCCESS_FILE + "`").
			SelectAll()
		if err != nil {
			return
		}

		for rows.Next() {
			var id string
			err = rows.Scan(&id)
			self.Success.old[id] = true
		}

	default:
		f, err := os.Open(SUCCESS_FILE_FULL)
		if err != nil {
			return
		}
		defer f.Close()
		b, _ := ioutil.ReadAll(f)
		b[0] = '{'
		json.Unmarshal(append(b, '}'), &self.Success.old)
	}
	logs.Log.Informational(" *     读出 %v 条成功记录\n", len(self.Success.old))
}

// 读取失败记录
func (self *History) ReadFailure(provider string, inherit bool) {
	self.RWMutex.Lock()
	self.provider = provider
	self.RWMutex.Unlock()

	if !inherit {
		// 不继承历史记录时
		self.Failure.old = make(map[string]map[string]bool)
		self.Failure.new = make(map[string]map[string]bool)
		self.Failure.inheritable = false
		return

	} else if self.Failure.inheritable {
		// 本次与上次均继承历史记录时
		return

	} else {
		// 上次没有继承历史记录，但本次继承时
		self.Failure.old = make(map[string]map[string]bool)
		self.Failure.new = make(map[string]map[string]bool)
		self.Failure.inheritable = true
	}
	var fLen int
	switch provider {
	case "mgo":
		var docs = []interface{}{}
		s, c, err := mgo.Open(MGO_DB, FAILURE_FILE)
		if err != nil {
			logs.Log.Error("从mgo读取成功记录: %v", err)
			return
		}
		c.Find(nil).All(&docs)
		mgo.Close(s)

		for _, v := range docs {
			key := v.(bson.M)["_id"].(string)
			req, err := context.UnSerialize(key)
			if err != nil {
				continue
			}
			spName := req.GetSpiderName()
			if _, ok := self.Failure.old[spName]; !ok {
				self.Failure.old[spName] = make(map[string]bool)
				self.Failure.new[spName] = make(map[string]bool)
			}
			self.Failure.old[spName][key] = true
			fLen++
		}

	case "mysql":
		db, ok := mysql.MysqlPool.GetOne().(*mysql.MysqlSrc)
		if !ok || db == nil {
			logs.Log.Error("链接Mysql数据库超时，无法读取成功记录！")
			return
		}
		rows, err := mysql.New(db.DB).
			SetTableName("`" + FAILURE_FILE + "`").
			SelectAll()
		if err != nil {
			// logs.Log.Error("读取Mysql数据库中成功记录失败：%v", err)
			return
		}

		mysql.MysqlPool.Free(db)

		for rows.Next() {
			var id int
			var failure string
			err = rows.Scan(&id, &failure)
			logs.Log.Error("Scan   %v", err)
			logs.Log.Error("failure   %v", failure)
			req, err := context.UnSerialize(failure)
			if err != nil {
				logs.Log.Error("UnSerialize   %v", err)
				continue
			}
			spName := req.GetSpiderName()
			if _, ok := self.Failure.old[spName]; !ok {
				self.Failure.old[spName] = make(map[string]bool)
				self.Failure.new[spName] = make(map[string]bool)
			}
			self.Failure.old[spName][failure] = true
			fLen++
		}

	default:
		f, err := os.Open(FAILURE_FILE_FULL)
		if err != nil {
			return
		}
		b, _ := ioutil.ReadAll(f)
		b[0] = '{'
		json.Unmarshal(
			append(b, '}'),
			&self.Failure.old,
		)
		for k, v := range self.Failure.old {
			self.Failure.new[k] = make(map[string]bool)
			fLen += len(v)
		}
		f.Close()
	}
	logs.Log.Informational(" *     读出 %v 条失败记录\n", fLen)
}

// 清空缓存，但不输出
func (self *History) Empty() {
	self.RWMutex.Lock()
	self.Success.new = make(map[string]bool)
	self.Success.old = make(map[string]bool)
	self.Failure.new = make(map[string]map[string]bool)
	self.Failure.old = make(map[string]map[string]bool)
	self.RWMutex.Unlock()
}

// I/O输出成功记录，但不清缓存
func (self *History) FlushSuccess(provider string) {
	self.RWMutex.Lock()
	self.provider = provider
	self.RWMutex.Unlock()
	sucLen := self.Success.flush(provider)
	logs.Log.Informational(" *     新增 %v 条成功记录\n", sucLen)
}

// I/O输出失败记录，但不清缓存
func (self *History) FlushFailure(provider string) {
	self.RWMutex.Lock()
	self.provider = provider
	self.RWMutex.Unlock()
	failLen := self.Failure.flush(provider)
	logs.Log.Informational(" *     新增 %v 条失败记录\n", failLen)
}

var once = new(sync.Once)

func mkdir() {
	p, _ := path.Split(SUCCESS_FILE_FULL)
	// 创建/打开目录
	d, err := os.Stat(p)
	if err != nil || !d.IsDir() {
		if err = os.MkdirAll(p, 0777); err != nil {
			logs.Log.Error("Error: %v\n", err)
		}
	}

	p, _ = path.Split(FAILURE_FILE_FULL)
	// 创建/打开目录
	d, err = os.Stat(p)
	if err != nil || !d.IsDir() {
		if err = os.MkdirAll(p, 0777); err != nil {
			logs.Log.Error("Error: %v\n", err)
		}
	}
}
<commit_msg>历史记录模块修复bug<commit_after>package history

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"path"
	"sync"

	"gopkg.in/mgo.v2/bson"

	"github.com/henrylee2cn/pholcus/app/downloader/context"
	"github.com/henrylee2cn/pholcus/common/mgo"
	"github.com/henrylee2cn/pholcus/common/mysql"
	"github.com/henrylee2cn/pholcus/config"
	"github.com/henrylee2cn/pholcus/logs"
)

type Historier interface {
	// 读取成功记录
	ReadSuccess(provider string, inherit bool)
	// 更新或加入成功记录
	UpsertSuccess(Record) bool
	// 删除成功记录
	DeleteSuccess(Record)
	// I/O输出成功记录，但不清缓存
	FlushSuccess(provider string)

	// 读取失败记录
	ReadFailure(provider string, inherit bool)
	// 更新或加入失败记录
	UpsertFailure(*context.Request) bool
	// 删除失败记录
	DeleteFailure(*context.Request)
	// I/O输出失败记录，但不清缓存
	FlushFailure(provider string)
	// 获取指定蜘蛛在上一次运行时失败的请求
	PullFailure(spiderName string) []*context.Request

	// 清空缓存，但不输出
	Empty()
}

type Record interface {
	GetUrl() string
	GetMethod() string
}

var (
	MGO_DB = config.MGO.DB

	SUCCESS_FILE = config.HISTORY.FILE_NAME + "_y"
	FAILURE_FILE = config.HISTORY.FILE_NAME + "_n"

	SUCCESS_FILE_FULL = config.COMM_PATH.CACHE + "/" + SUCCESS_FILE
	FAILURE_FILE_FULL = config.COMM_PATH.CACHE + "/" + FAILURE_FILE
)

type History struct {
	*Success
	*Failure
	provider string
	sync.RWMutex
}

func New() Historier {
	return &History{
		Success: &Success{
			new: make(map[string]bool),
			old: make(map[string]bool),
		},
		Failure: &Failure{
			new: make(map[string]map[string]bool),
			old: make(map[string]map[string]bool),
		},
	}
}

// 读取成功记录
func (self *History) ReadSuccess(provider string, inherit bool) {
	self.RWMutex.Lock()
	self.provider = provider
	self.RWMutex.Unlock()

	if !inherit {
		// 不继承历史记录时
		self.Success.old = make(map[string]bool)
		self.Success.new = make(map[string]bool)
		self.Success.inheritable = false
		return

	} else if self.Success.inheritable {
		// 本次与上次均继承历史记录时
		return

	} else {
		// 上次没有继承历史记录，但本次继承时
		self.Success.old = make(map[string]bool)
		self.Success.new = make(map[string]bool)
		self.Success.inheritable = true
	}

	switch provider {
	case "mgo":
		var docs = map[string]interface{}{}
		err := mgo.Mgo(&docs, "find", map[string]interface{}{
			"Database":   MGO_DB,
			"Collection": SUCCESS_FILE,
		})
		if err != nil {
			logs.Log.Error("从mgo读取成功记录: %v", err)
			return
		}
		for _, v := range docs["Docs"].([]interface{}) {
			self.Success.old[v.(bson.M)["_id"].(string)] = true
		}

	case "mysql":
		db, ok := mysql.MysqlPool.GetOne().(*mysql.MysqlSrc)
		if !ok || db == nil {
			// logs.Log.Error("链接Mysql数据库超时，无法读取成功记录！")
			return
		}
		defer mysql.MysqlPool.Free(db)
		rows, err := mysql.New(db.DB).
			SetTableName("`" + SUCCESS_FILE + "`").
			SelectAll()
		if err != nil {
			return
		}

		for rows.Next() {
			var id string
			err = rows.Scan(&id)
			self.Success.old[id] = true
		}

	default:
		f, err := os.Open(SUCCESS_FILE_FULL)
		if err != nil {
			return
		}
		defer f.Close()
		b, _ := ioutil.ReadAll(f)
		b[0] = '{'
		json.Unmarshal(append(b, '}'), &self.Success.old)
	}
	logs.Log.Informational(" *     读出 %v 条成功记录\n", len(self.Success.old))
}

// 读取失败记录
func (self *History) ReadFailure(provider string, inherit bool) {
	self.RWMutex.Lock()
	self.provider = provider
	self.RWMutex.Unlock()

	if !inherit {
		// 不继承历史记录时
		self.Failure.old = make(map[string]map[string]bool)
		self.Failure.new = make(map[string]map[string]bool)
		self.Failure.inheritable = false
		return

	} else if self.Failure.inheritable {
		// 本次与上次均继承历史记录时
		return

	} else {
		// 上次没有继承历史记录，但本次继承时
		self.Failure.old = make(map[string]map[string]bool)
		self.Failure.new = make(map[string]map[string]bool)
		self.Failure.inheritable = true
	}
	var fLen int
	switch provider {
	case "mgo":
		var docs = []interface{}{}
		s, c, err := mgo.Open(MGO_DB, FAILURE_FILE)
		if err != nil {
			logs.Log.Error("从mgo读取成功记录: %v", err)
			return
		}
		c.Find(nil).All(&docs)
		mgo.Close(s)

		for _, v := range docs {
			key := v.(bson.M)["_id"].(string)
			req, err := context.UnSerialize(key)
			if err != nil {
				continue
			}
			spName := req.GetSpiderName()
			if _, ok := self.Failure.old[spName]; !ok {
				self.Failure.old[spName] = make(map[string]bool)
				self.Failure.new[spName] = make(map[string]bool)
			}
			self.Failure.old[spName][key] = true
			fLen++
		}

	case "mysql":
		db, ok := mysql.MysqlPool.GetOne().(*mysql.MysqlSrc)
		if !ok || db == nil {
			logs.Log.Error("链接Mysql数据库超时，无法读取成功记录！")
			return
		}
		rows, err := mysql.New(db.DB).
			SetTableName("`" + FAILURE_FILE + "`").
			SelectAll()
		if err != nil {
			// logs.Log.Error("读取Mysql数据库中成功记录失败：%v", err)
			return
		}

		mysql.MysqlPool.Free(db)

		for rows.Next() {
			var id int
			var failure string
			err = rows.Scan(&id, &failure)
			req, err := context.UnSerialize(failure)
			if err != nil {
				continue
			}
			spName := req.GetSpiderName()
			if _, ok := self.Failure.old[spName]; !ok {
				self.Failure.old[spName] = make(map[string]bool)
				self.Failure.new[spName] = make(map[string]bool)
			}
			self.Failure.old[spName][failure] = true
			fLen++
		}

	default:
		f, err := os.Open(FAILURE_FILE_FULL)
		if err != nil {
			return
		}
		b, _ := ioutil.ReadAll(f)
		b[0] = '{'
		json.Unmarshal(
			append(b, '}'),
			&self.Failure.old,
		)
		for k, v := range self.Failure.old {
			self.Failure.new[k] = make(map[string]bool)
			fLen += len(v)
		}
		f.Close()
	}
	logs.Log.Informational(" *     读出 %v 条失败记录\n", fLen)
}

// 清空缓存，但不输出
func (self *History) Empty() {
	self.RWMutex.Lock()
	self.Success.new = make(map[string]bool)
	self.Success.old = make(map[string]bool)
	self.Failure.new = make(map[string]map[string]bool)
	self.Failure.old = make(map[string]map[string]bool)
	self.RWMutex.Unlock()
}

// I/O输出成功记录，但不清缓存
func (self *History) FlushSuccess(provider string) {
	self.RWMutex.Lock()
	self.provider = provider
	self.RWMutex.Unlock()
	sucLen := self.Success.flush(provider)
	logs.Log.Informational(" *     新增 %v 条成功记录\n", sucLen)
}

// I/O输出失败记录，但不清缓存
func (self *History) FlushFailure(provider string) {
	self.RWMutex.Lock()
	self.provider = provider
	self.RWMutex.Unlock()
	failLen := self.Failure.flush(provider)
	logs.Log.Informational(" *     新增 %v 条失败记录\n", failLen)
}

var once = new(sync.Once)

func mkdir() {
	p, _ := path.Split(SUCCESS_FILE_FULL)
	// 创建/打开目录
	d, err := os.Stat(p)
	if err != nil || !d.IsDir() {
		if err = os.MkdirAll(p, 0777); err != nil {
			logs.Log.Error("Error: %v\n", err)
		}
	}

	p, _ = path.Split(FAILURE_FILE_FULL)
	// 创建/打开目录
	d, err = os.Stat(p)
	if err != nil || !d.IsDir() {
		if err = os.MkdirAll(p, 0777); err != nil {
			logs.Log.Error("Error: %v\n", err)
		}
	}
}
<|endoftext|>


<commit_before>package download

import "testing"

func TestClean(t *testing.T) {
	cases := []struct {
		path string
		want string
	}{
		{"github.com/foo/bar", "github.com/foo/bar"},
		{"https://github.com/foo/bar", "github.com/foo/bar"},
		{"https://user@github.com/foo/bar", "github.com/foo/bar"},
		{"github.com/foo/bar/v2", "github.com/foo/bar/v2"},
	}

	for _, tt := range cases {
		got, err := Clean(tt.path)
		if err != nil {
			t.Fatal(err)
		}

		if got != tt.want {
			t.Errorf("Clean(%q) = %q, want %q", tt.path, got, tt.want)
		}
	}
}
<commit_msg>Add more tests.<commit_after>package download

import "testing"

func TestClean(t *testing.T) {
	cases := []struct {
		path string
		want string
	}{
		{"github.com/foo/bar", "github.com/foo/bar"},
		{"https://github.com/foo/bar", "github.com/foo/bar"},
		{"https://user@github.com/foo/bar", "github.com/foo/bar"},
		{"github.com/foo/bar/v2", "github.com/foo/bar/v2"},
		{"https://github.com/foo/bar/v2", "github.com/foo/bar/v2"},
		{"https://user@github.com/foo/bar/v2", "github.com/foo/bar/v2"},
	}

	for _, tt := range cases {
		got := Clean(tt.path)

		if got != tt.want {
			t.Errorf("Clean(%q) = %q, want %q", tt.path, got, tt.want)
		}
	}
}
<|endoftext|>


<commit_before>/* Collection and storage management. */
package db

import (
	"encoding/json"
	"fmt"
	"github.com/HouzuoGuo/tiedot/tdlog"
	"io"
	"io/ioutil"
	"math/rand"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	PART_NUM_FILE      = "number_of_partitions" // A database configuration file's name, content is a single number
	AUTO_SYNC_INTERVAL = 2000                   // Data file auto-save interval in milliseconds (do not set too small)
)

// Database structures.
type DB struct {
	path         string          // root path of database directory
	numParts     int             // total number of partitions
	cols         map[string]*Col // collection by name lookup
	schemaLock   *sync.RWMutex   // Schema-related operations use ExLock, other operations use RLock
	autoSync     *time.Ticker    // Synchronize all data files at regular interval
	autoSyncStop chan struct{}   // Inform auto-synchronize routine to stop (on DB shutdown)
}

// Open database and load all collections & indexes.
func OpenDB(dbPath string) (*DB, error) {
	// New collection documents get their ID from this RNG
	rand.Seed(time.Now().UnixNano())
	db := &DB{path: dbPath, schemaLock: new(sync.RWMutex), autoSyncStop: make(chan struct{})}
	return db, db.load()
}

// Load all collection schema.
func (db *DB) load() error {
	// If opening an empty database, create the DB directory and PART_NUM_FILE.
	var numPartsAssumed = false
	numPartsFilePath := path.Join(db.path, PART_NUM_FILE)
	if err := os.MkdirAll(db.path, 0700); err != nil {
		return err
	}
	if partNumFile, err := os.Stat(numPartsFilePath); err != nil {
		// The new database has as many partitions as there are CPUs
		if err := ioutil.WriteFile(numPartsFilePath, []byte(strconv.Itoa(runtime.NumCPU())), 0600); err != nil {
			return err
		}
		numPartsAssumed = true
	} else if partNumFile.IsDir() {
		return fmt.Errorf("Database config file %s is actually a directory, is database path correct?", PART_NUM_FILE)
	}
	// Get number of partitions from the text file
	if numParts, err := ioutil.ReadFile(numPartsFilePath); err != nil {
		return err
	} else if db.numParts, err = strconv.Atoi(string(numParts)); err != nil {
		return err
	}
	// Look for collection directories
	db.cols = make(map[string]*Col)
	dirContent, err := ioutil.ReadDir(db.path)
	if err != nil {
		return err
	}
	for _, maybeColDir := range dirContent {
		if !maybeColDir.IsDir() {
			continue
		}
		if numPartsAssumed {
			return fmt.Errorf("Please manually repair database partition number config file %s", numPartsFilePath)
		}
		if db.cols[maybeColDir.Name()], err = OpenCol(db, maybeColDir.Name()); err != nil {
			return err
		}
	}
	// Initialize auto-sync (synchronize all data files regularly)
	if db.autoSync == nil {
		db.autoSync = time.NewTicker(AUTO_SYNC_INTERVAL * time.Millisecond)
		go func() {
			for {
				select {
				case <-db.autoSync.C:
					if err := db.Sync(); err == nil {
						tdlog.Printf("Background Auto-Sync on %s: OK", db.path)
					} else {
						tdlog.Errorf("Background Auto-Sync on %s: Failed with error: %v", db.path, err)
					}
				case <-db.autoSyncStop:
					db.autoSync.Stop()
					return
				}
			}
		}()
	}
	return err
}

// Synchronize all data files to disk.
func (db *DB) Sync() error {
	db.schemaLock.Lock()
	defer db.schemaLock.Unlock()
	errs := make([]error, 0, 0)
	for _, col := range db.cols {
		if err := col.Sync(); err != nil {
			errs = append(errs, err)
		}
	}
	if len(errs) == 0 {
		return nil
	}
	return fmt.Errorf("%v", errs)
}

// Close all database files. Do not use the DB afterwards!
func (db *DB) Close() error {
	db.schemaLock.Lock()
	defer db.schemaLock.Unlock()
	close(db.autoSyncStop)
	errs := make([]error, 0, 0)
	for _, col := range db.cols {
		if err := col.Close(); err != nil {
			errs = append(errs, err)
		}
	}
	if len(errs) == 0 {
		return nil
	}
	return fmt.Errorf("%v", errs)
}

// Create a new collection.
func (db *DB) Create(name string) error {
	db.schemaLock.Lock()
	defer db.schemaLock.Unlock()
	if _, exists := db.cols[name]; exists {
		return fmt.Errorf("Collection %s already exists", name)
	} else if err := os.MkdirAll(path.Join(db.path, name), 0700); err != nil {
		return err
	} else if db.cols[name], err = OpenCol(db, name); err != nil {
		return err
	}
	return nil
}

// Return all collection names.
func (db *DB) AllCols() (ret []string) {
	db.schemaLock.RLock()
	defer db.schemaLock.RUnlock()
	ret = make([]string, 0, len(db.cols))
	for name, _ := range db.cols {
		ret = append(ret, name)
	}
	return
}

// Use the return value to interact with collection. Return value may be nil if the collection does not exist.
func (db *DB) Use(name string) *Col {
	db.schemaLock.RLock()
	defer db.schemaLock.RUnlock()
	if col, exists := db.cols[name]; exists {
		return col
	}
	return nil
}

// Rename a collection.
func (db *DB) Rename(oldName, newName string) error {
	db.schemaLock.Lock()
	defer db.schemaLock.Unlock()
	if _, exists := db.cols[oldName]; !exists {
		return fmt.Errorf("Collection %s does not exist", oldName)
	} else if _, exists := db.cols[newName]; exists {
		return fmt.Errorf("Collection %s already exists", newName)
	} else if newName == oldName {
		return fmt.Errorf("Old and new names are the same")
	} else if err := db.cols[oldName].Close(); err != nil {
		return err
	} else if err := os.Rename(path.Join(db.path, oldName), path.Join(db.path, newName)); err != nil {
		return err
	} else if db.cols[newName], err = OpenCol(db, newName); err != nil {
		return err
	}
	delete(db.cols, oldName)
	return nil
}

// Truncate a collection - delete all documents and clear
func (db *DB) Truncate(name string) error {
	db.schemaLock.Lock()
	defer db.schemaLock.Unlock()
	if _, exists := db.cols[name]; !exists {
		return fmt.Errorf("Collection %s does not exist", name)
	} else if err := db.cols[name].Sync(); err != nil {
		return err
	}
	col := db.cols[name]
	for i := 0; i < db.numParts; i++ {
		if err := col.parts[i].Clear(); err != nil {
			return err
		}
		for _, ht := range col.hts[i] {
			if err := ht.Clear(); err != nil {
				return err
			}
		}
	}
	return nil
}

// Scrub a collection - fix corrupted documents and de-fragment free space.
func (db *DB) Scrub(name string) error {
	db.schemaLock.Lock()
	defer db.schemaLock.Unlock()
	if _, exists := db.cols[name]; !exists {
		return fmt.Errorf("Collection %s does not exist", name)
	} else if err := db.cols[name].Sync(); err != nil {
		return err
	}
	// Prepare a temporary collection in file system
	tmpColName := fmt.Sprintf("scrub-%s-%d", name, time.Now().UnixNano())
	tmpColDir := path.Join(db.path, tmpColName)
	if err := os.MkdirAll(tmpColDir, 0700); err != nil {
		return err
	}
	// Mirror indexes from original collection
	for _, idxPath := range db.cols[name].indexPaths {
		if err := os.MkdirAll(path.Join(tmpColDir, strings.Join(idxPath, INDEX_PATH_SEP)), 0700); err != nil {
			return err
		}
	}
	// Iterate through all documents and put them into the temporary collection
	tmpCol, err := OpenCol(db, tmpColName)
	if err != nil {
		return err
	}
	db.cols[name].ForEachDoc(false, func(id int, doc []byte) (moveOn bool) {
		var docObj map[string]interface{}
		if err := json.Unmarshal([]byte(doc), &docObj); err != nil {
			// Skip corrupted document
			return true
		}
		if err := tmpCol.InsertRecovery(id, docObj); err != nil {
			tdlog.Errorf("Scrub %s: failed to insert back document %v", name, docObj)
		}
		return true
	})
	if err := tmpCol.Close(); err != nil {
		return err
	}
	// Replace the original collection with the "temporary" one
	db.cols[name].Close()
	if err := os.RemoveAll(path.Join(db.path, name)); err != nil {
		return err
	}
	if err := os.Rename(path.Join(db.path, tmpColName), path.Join(db.path, name)); err != nil {
		return err
	}
	if db.cols[name], err = OpenCol(db, name); err != nil {
		return err
	}
	return nil
}

// Drop a collection and lose all of its documents and indexes.
func (db *DB) Drop(name string) error {
	db.schemaLock.Lock()
	defer db.schemaLock.Unlock()
	if _, exists := db.cols[name]; !exists {
		return fmt.Errorf("Collection %s does not exist", name)
	} else if err := db.cols[name].Close(); err != nil {
		return err
	} else if err := os.RemoveAll(path.Join(db.path, name)); err != nil {
		return err
	}
	delete(db.cols, name)
	return nil
}

// Copy this database into destination directory.
func (db *DB) Dump(dest string) error {
	db.schemaLock.Lock()
	defer db.schemaLock.Unlock()
	for _, col := range db.cols {
		if err := col.Sync(); err != nil {
			return err
		}
	}
	cpFun := func(currPath string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() {
			relPath, err := filepath.Rel(db.path, currPath)
			if err != nil {
				return err
			}
			destDir := path.Join(dest, relPath)
			if err := os.MkdirAll(destDir, 0700); err != nil {
				return err
			}
			tdlog.Printf("Dump: created directory %s", destDir)
		} else {
			src, err := os.Open(currPath)
			if err != nil {
				return err
			}
			relPath, err := filepath.Rel(db.path, currPath)
			if err != nil {
				return err
			}
			destPath := path.Join(dest, relPath)
			if _, fileExists := os.Open(destPath); fileExists == nil {
				return fmt.Errorf("Destination file %s already exists", destPath)
			}
			destFile, err := os.Create(destPath)
			if err != nil {
				return err
			}
			written, err := io.Copy(destFile, src)
			if err != nil {
				return err
			}
			tdlog.Printf("Dump: copied file %s, size is %d", destPath, written)
		}
		return nil
	}
	return filepath.Walk(db.path, cpFun)
}
<commit_msg>trim num_partitions file content from spaces and newline<commit_after>/* Collection and storage management. */
package db

import (
	"encoding/json"
	"fmt"
	"github.com/HouzuoGuo/tiedot/tdlog"
	"io"
	"io/ioutil"
	"math/rand"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	PART_NUM_FILE      = "number_of_partitions" // A database configuration file's name, content is a single number
	AUTO_SYNC_INTERVAL = 2000                   // Data file auto-save interval in milliseconds (do not set too small)
)

// Database structures.
type DB struct {
	path         string          // root path of database directory
	numParts     int             // total number of partitions
	cols         map[string]*Col // collection by name lookup
	schemaLock   *sync.RWMutex   // Schema-related operations use ExLock, other operations use RLock
	autoSync     *time.Ticker    // Synchronize all data files at regular interval
	autoSyncStop chan struct{}   // Inform auto-synchronize routine to stop (on DB shutdown)
}

// Open database and load all collections & indexes.
func OpenDB(dbPath string) (*DB, error) {
	// New collection documents get their ID from this RNG
	rand.Seed(time.Now().UnixNano())
	db := &DB{path: dbPath, schemaLock: new(sync.RWMutex), autoSyncStop: make(chan struct{})}
	return db, db.load()
}

// Load all collection schema.
func (db *DB) load() error {
	// If opening an empty database, create the DB directory and PART_NUM_FILE.
	var numPartsAssumed = false
	numPartsFilePath := path.Join(db.path, PART_NUM_FILE)
	if err := os.MkdirAll(db.path, 0700); err != nil {
		return err
	}
	if partNumFile, err := os.Stat(numPartsFilePath); err != nil {
		// The new database has as many partitions as there are CPUs
		if err := ioutil.WriteFile(numPartsFilePath, []byte(strconv.Itoa(runtime.NumCPU())), 0600); err != nil {
			return err
		}
		numPartsAssumed = true
	} else if partNumFile.IsDir() {
		return fmt.Errorf("Database config file %s is actually a directory, is database path correct?", PART_NUM_FILE)
	}
	// Get number of partitions from the text file
	if numParts, err := ioutil.ReadFile(numPartsFilePath); err != nil {
		return err
	} else if db.numParts, err = strconv.Atoi(strings.Trim(string(numParts), "\r\n ")); err != nil {
		return err
	}
	// Look for collection directories
	db.cols = make(map[string]*Col)
	dirContent, err := ioutil.ReadDir(db.path)
	if err != nil {
		return err
	}
	for _, maybeColDir := range dirContent {
		if !maybeColDir.IsDir() {
			continue
		}
		if numPartsAssumed {
			return fmt.Errorf("Please manually repair database partition number config file %s", numPartsFilePath)
		}
		if db.cols[maybeColDir.Name()], err = OpenCol(db, maybeColDir.Name()); err != nil {
			return err
		}
	}
	// Initialize auto-sync (synchronize all data files regularly)
	if db.autoSync == nil {
		db.autoSync = time.NewTicker(AUTO_SYNC_INTERVAL * time.Millisecond)
		go func() {
			for {
				select {
				case <-db.autoSync.C:
					if err := db.Sync(); err == nil {
						tdlog.Printf("Background Auto-Sync on %s: OK", db.path)
					} else {
						tdlog.Errorf("Background Auto-Sync on %s: Failed with error: %v", db.path, err)
					}
				case <-db.autoSyncStop:
					db.autoSync.Stop()
					return
				}
			}
		}()
	}
	return err
}

// Synchronize all data files to disk.
func (db *DB) Sync() error {
	db.schemaLock.Lock()
	defer db.schemaLock.Unlock()
	errs := make([]error, 0, 0)
	for _, col := range db.cols {
		if err := col.Sync(); err != nil {
			errs = append(errs, err)
		}
	}
	if len(errs) == 0 {
		return nil
	}
	return fmt.Errorf("%v", errs)
}

// Close all database files. Do not use the DB afterwards!
func (db *DB) Close() error {
	db.schemaLock.Lock()
	defer db.schemaLock.Unlock()
	close(db.autoSyncStop)
	errs := make([]error, 0, 0)
	for _, col := range db.cols {
		if err := col.Close(); err != nil {
			errs = append(errs, err)
		}
	}
	if len(errs) == 0 {
		return nil
	}
	return fmt.Errorf("%v", errs)
}

// Create a new collection.
func (db *DB) Create(name string) error {
	db.schemaLock.Lock()
	defer db.schemaLock.Unlock()
	if _, exists := db.cols[name]; exists {
		return fmt.Errorf("Collection %s already exists", name)
	} else if err := os.MkdirAll(path.Join(db.path, name), 0700); err != nil {
		return err
	} else if db.cols[name], err = OpenCol(db, name); err != nil {
		return err
	}
	return nil
}

// Return all collection names.
func (db *DB) AllCols() (ret []string) {
	db.schemaLock.RLock()
	defer db.schemaLock.RUnlock()
	ret = make([]string, 0, len(db.cols))
	for name, _ := range db.cols {
		ret = append(ret, name)
	}
	return
}

// Use the return value to interact with collection. Return value may be nil if the collection does not exist.
func (db *DB) Use(name string) *Col {
	db.schemaLock.RLock()
	defer db.schemaLock.RUnlock()
	if col, exists := db.cols[name]; exists {
		return col
	}
	return nil
}

// Rename a collection.
func (db *DB) Rename(oldName, newName string) error {
	db.schemaLock.Lock()
	defer db.schemaLock.Unlock()
	if _, exists := db.cols[oldName]; !exists {
		return fmt.Errorf("Collection %s does not exist", oldName)
	} else if _, exists := db.cols[newName]; exists {
		return fmt.Errorf("Collection %s already exists", newName)
	} else if newName == oldName {
		return fmt.Errorf("Old and new names are the same")
	} else if err := db.cols[oldName].Close(); err != nil {
		return err
	} else if err := os.Rename(path.Join(db.path, oldName), path.Join(db.path, newName)); err != nil {
		return err
	} else if db.cols[newName], err = OpenCol(db, newName); err != nil {
		return err
	}
	delete(db.cols, oldName)
	return nil
}

// Truncate a collection - delete all documents and clear
func (db *DB) Truncate(name string) error {
	db.schemaLock.Lock()
	defer db.schemaLock.Unlock()
	if _, exists := db.cols[name]; !exists {
		return fmt.Errorf("Collection %s does not exist", name)
	} else if err := db.cols[name].Sync(); err != nil {
		return err
	}
	col := db.cols[name]
	for i := 0; i < db.numParts; i++ {
		if err := col.parts[i].Clear(); err != nil {
			return err
		}
		for _, ht := range col.hts[i] {
			if err := ht.Clear(); err != nil {
				return err
			}
		}
	}
	return nil
}

// Scrub a collection - fix corrupted documents and de-fragment free space.
func (db *DB) Scrub(name string) error {
	db.schemaLock.Lock()
	defer db.schemaLock.Unlock()
	if _, exists := db.cols[name]; !exists {
		return fmt.Errorf("Collection %s does not exist", name)
	} else if err := db.cols[name].Sync(); err != nil {
		return err
	}
	// Prepare a temporary collection in file system
	tmpColName := fmt.Sprintf("scrub-%s-%d", name, time.Now().UnixNano())
	tmpColDir := path.Join(db.path, tmpColName)
	if err := os.MkdirAll(tmpColDir, 0700); err != nil {
		return err
	}
	// Mirror indexes from original collection
	for _, idxPath := range db.cols[name].indexPaths {
		if err := os.MkdirAll(path.Join(tmpColDir, strings.Join(idxPath, INDEX_PATH_SEP)), 0700); err != nil {
			return err
		}
	}
	// Iterate through all documents and put them into the temporary collection
	tmpCol, err := OpenCol(db, tmpColName)
	if err != nil {
		return err
	}
	db.cols[name].ForEachDoc(false, func(id int, doc []byte) (moveOn bool) {
		var docObj map[string]interface{}
		if err := json.Unmarshal([]byte(doc), &docObj); err != nil {
			// Skip corrupted document
			return true
		}
		if err := tmpCol.InsertRecovery(id, docObj); err != nil {
			tdlog.Errorf("Scrub %s: failed to insert back document %v", name, docObj)
		}
		return true
	})
	if err := tmpCol.Close(); err != nil {
		return err
	}
	// Replace the original collection with the "temporary" one
	db.cols[name].Close()
	if err := os.RemoveAll(path.Join(db.path, name)); err != nil {
		return err
	}
	if err := os.Rename(path.Join(db.path, tmpColName), path.Join(db.path, name)); err != nil {
		return err
	}
	if db.cols[name], err = OpenCol(db, name); err != nil {
		return err
	}
	return nil
}

// Drop a collection and lose all of its documents and indexes.
func (db *DB) Drop(name string) error {
	db.schemaLock.Lock()
	defer db.schemaLock.Unlock()
	if _, exists := db.cols[name]; !exists {
		return fmt.Errorf("Collection %s does not exist", name)
	} else if err := db.cols[name].Close(); err != nil {
		return err
	} else if err := os.RemoveAll(path.Join(db.path, name)); err != nil {
		return err
	}
	delete(db.cols, name)
	return nil
}

// Copy this database into destination directory.
func (db *DB) Dump(dest string) error {
	db.schemaLock.Lock()
	defer db.schemaLock.Unlock()
	for _, col := range db.cols {
		if err := col.Sync(); err != nil {
			return err
		}
	}
	cpFun := func(currPath string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() {
			relPath, err := filepath.Rel(db.path, currPath)
			if err != nil {
				return err
			}
			destDir := path.Join(dest, relPath)
			if err := os.MkdirAll(destDir, 0700); err != nil {
				return err
			}
			tdlog.Printf("Dump: created directory %s", destDir)
		} else {
			src, err := os.Open(currPath)
			if err != nil {
				return err
			}
			relPath, err := filepath.Rel(db.path, currPath)
			if err != nil {
				return err
			}
			destPath := path.Join(dest, relPath)
			if _, fileExists := os.Open(destPath); fileExists == nil {
				return fmt.Errorf("Destination file %s already exists", destPath)
			}
			destFile, err := os.Create(destPath)
			if err != nil {
				return err
			}
			written, err := io.Copy(destFile, src)
			if err != nil {
				return err
			}
			tdlog.Printf("Dump: copied file %s, size is %d", destPath, written)
		}
		return nil
	}
	return filepath.Walk(db.path, cpFun)
}
<|endoftext|>


<commit_before>package db

import (
	"database/sql/driver"
	"fmt"
	"io"
	"time"

	"github.com/mattn/go-sqlite3"
)

const bkDelay = 250

// Config represents the configuration of the SQLite database.
type Config struct {
	DSN    string // Datasource name
	Memory bool   // In-memory database enabled?
}

// NewConfig returns an instance of Config for the database at path.
func NewConfig() *Config {
	return &Config{}
}

// FQDSN returns the fully-qualified datasource name.
func (c *Config) FQDSN(path string) string {
	if c.DSN != "" {
		return fmt.Sprintf("file:%s?%s", path, c.DSN)
	}
	return path
}

// DB is the SQL database.
type DB struct {
	sqlite3conn *sqlite3.SQLiteConn // Driver connection to database.
	path        string              // Path to database file.
}

// Result represents the outcome of an operation that changes rows.
type Result struct {
	LastInsertID int64   `json:"last_insert_id,omitempty"`
	RowsAffected int64   `json:"rows_affected,omitempty"`
	Error        string  `json:"error,omitempty"`
	Time         float64 `json:"time,omitempty"`
}

// Rows represents the outcome of an operation that returns query data.
type Rows struct {
	Columns []string        `json:"columns,omitempty"`
	Types   []string        `json:"types,omitempty"`
	Values  [][]interface{} `json:"values,omitempty"`
	Error   string          `json:"error,omitempty"`
	Time    float64         `json:"time,omitempty"`
}

// Open an existing database, creating it if it does not exist.
func Open(dbPath string) (*DB, error) {
	d := sqlite3.SQLiteDriver{}
	dbc, err := d.Open(dbPath)
	if err != nil {
		return nil, err
	}

	return &DB{
		sqlite3conn: dbc.(*sqlite3.SQLiteConn),
		path:        dbPath,
	}, nil
}

// OpenWithConfiguration an existing database, creating it if it does not exist.
func OpenWithConfiguration(dbPath string, conf *Config) (*DB, error) {
	return Open(conf.FQDSN(dbPath))
}

// Close closes the underlying database connection.
func (db *DB) Close() error {
	return db.sqlite3conn.Close()
}

// Execute executes queries that modify the database.
func (db *DB) Execute(queries []string, tx, xTime bool) ([]*Result, error) {
	type Execer interface {
		Exec(query string, args []driver.Value) (driver.Result, error)
	}

	var allResults []*Result
	err := func() (error) {
		var execer Execer
		var rollback bool
		var t driver.Tx
		var err error

		// Check for the err, if set rollback.
		defer func() {
			if t != nil {
				if rollback {
					t.Rollback()
					return
				}
				t.Commit()
			}
		}()

		// handleError sets the error field on the given result. It returns
		// whether the caller should continue processing or break.
		handleError := func(result *Result, err error) bool {
			result.Error = err.Error()
			allResults = append(allResults, result)
			if tx {
				rollback = true // Will trigger the rollback.
				return false
			}
			return true
		}

		execer = db.sqlite3conn

		// Create the correct execution object, depending on whether a
		// transaction was requested.
		if tx {
			t, err = db.sqlite3conn.Begin()
			if err != nil {
				return err
			}
		}

		// Execute each query.
		for _, q := range queries {
			if q == "" {
				continue
			}

			result := &Result{}
			start := time.Now()

			r, err := execer.Exec(q, nil)
			if err != nil {
				if handleError(result, err) {
					continue
				}
				break
			}

			lid, err := r.LastInsertId()
			if err != nil {
				if handleError(result, err) {
					continue
				}
				break
			}
			result.LastInsertID = lid

			ra, err := r.RowsAffected()
			if err != nil {
				if handleError(result, err) {
					continue
				}
				break
			}
			result.RowsAffected = ra
			if xTime {
				result.Time = time.Now().Sub(start).Seconds()
			}
			allResults = append(allResults, result)
		}

		return nil
	}()

	return allResults, err
}

// Query executes queries that return rows, but don't modify the database.
func (db *DB) Query(queries []string, tx, xTime bool) ([]*Rows, error) {
	type Queryer interface {
		Query(query string, args []driver.Value) (driver.Rows, error)
	}

	var allRows []*Rows
	err := func() (err error) {
		var queryer Queryer
		var t driver.Tx
		defer func() {
			// XXX THIS DOESN'T ACTUALLY WORK! Might as WELL JUST COMMIT?
			if t != nil {
				if err != nil {
					t.Rollback()
					return
				}
				t.Commit()
			}
		}()

		queryer = db.sqlite3conn

		// Create the correct query object, depending on whether a
		// transaction was requested.
		if tx {
			t, err = db.sqlite3conn.Begin()
			if err != nil {
				return err
			}
		}

		for _, q := range queries {
			if q == "" {
				continue
			}

			rows := &Rows{}
			start := time.Now()

			rs, err := queryer.Query(q, nil)
			if err != nil {
				rows.Error = err.Error()
				allRows = append(allRows, rows)
				continue
			}
			defer rs.Close() // This adds to all defers, right? Nothing leaks? XXX Could consume memory. Perhaps anon would be best.
			columns := rs.Columns()

			rows.Columns = columns
			rows.Types = rs.(*sqlite3.SQLiteRows).DeclTypes()
			dest := make([]driver.Value, len(rows.Columns))
			for {
				err := rs.Next(dest)

				if err != nil {
					if err != io.EOF {
						rows.Error = err.Error()
					}
					break
				}

				values := make([]interface{}, len(rows.Columns))

				// Special case -- convert []uint8 to string. Perhaps this should be a config option.
				for i, v := range dest {
					switch u := v.(type) {
					case []uint8:
						values[i] = string(u)
					default:
						values[i] = u
					}
				}
				rows.Values = append(rows.Values, values)
			}
			if xTime {
				rows.Time = time.Now().Sub(start).Seconds()
			}
			allRows = append(allRows, rows)
		}

		return nil
	}()

	return allRows, err
}

// Backup writes a consistent snapshot of the database to the given file.
func (db *DB) Backup(path string) error {
	dstDB, err := Open(path)
	if err != nil {
		return err
	}

	bk, err := dstDB.sqlite3conn.Backup("main", db.sqlite3conn, "main")
	if err != nil {
		return err
	}

	for {
		done, err := bk.Step(-1)
		if err != nil {
			bk.Finish()
			return err
		}
		if done {
			break
		}
		time.Sleep(bkDelay * time.Millisecond)
	}

	if err := bk.Finish(); err != nil {
		return err
	}
	if err := dstDB.Close(); err != nil {
		return err
	}

	return nil
}
<commit_msg>'go fmt' fixes<commit_after>package db

import (
	"database/sql/driver"
	"fmt"
	"io"
	"time"

	"github.com/mattn/go-sqlite3"
)

const bkDelay = 250

// Config represents the configuration of the SQLite database.
type Config struct {
	DSN    string // Datasource name
	Memory bool   // In-memory database enabled?
}

// NewConfig returns an instance of Config for the database at path.
func NewConfig() *Config {
	return &Config{}
}

// FQDSN returns the fully-qualified datasource name.
func (c *Config) FQDSN(path string) string {
	if c.DSN != "" {
		return fmt.Sprintf("file:%s?%s", path, c.DSN)
	}
	return path
}

// DB is the SQL database.
type DB struct {
	sqlite3conn *sqlite3.SQLiteConn // Driver connection to database.
	path        string              // Path to database file.
}

// Result represents the outcome of an operation that changes rows.
type Result struct {
	LastInsertID int64   `json:"last_insert_id,omitempty"`
	RowsAffected int64   `json:"rows_affected,omitempty"`
	Error        string  `json:"error,omitempty"`
	Time         float64 `json:"time,omitempty"`
}

// Rows represents the outcome of an operation that returns query data.
type Rows struct {
	Columns []string        `json:"columns,omitempty"`
	Types   []string        `json:"types,omitempty"`
	Values  [][]interface{} `json:"values,omitempty"`
	Error   string          `json:"error,omitempty"`
	Time    float64         `json:"time,omitempty"`
}

// Open an existing database, creating it if it does not exist.
func Open(dbPath string) (*DB, error) {
	d := sqlite3.SQLiteDriver{}
	dbc, err := d.Open(dbPath)
	if err != nil {
		return nil, err
	}

	return &DB{
		sqlite3conn: dbc.(*sqlite3.SQLiteConn),
		path:        dbPath,
	}, nil
}

// OpenWithConfiguration an existing database, creating it if it does not exist.
func OpenWithConfiguration(dbPath string, conf *Config) (*DB, error) {
	return Open(conf.FQDSN(dbPath))
}

// Close closes the underlying database connection.
func (db *DB) Close() error {
	return db.sqlite3conn.Close()
}

// Execute executes queries that modify the database.
func (db *DB) Execute(queries []string, tx, xTime bool) ([]*Result, error) {
	type Execer interface {
		Exec(query string, args []driver.Value) (driver.Result, error)
	}

	var allResults []*Result
	err := func() error {
		var execer Execer
		var rollback bool
		var t driver.Tx
		var err error

		// Check for the err, if set rollback.
		defer func() {
			if t != nil {
				if rollback {
					t.Rollback()
					return
				}
				t.Commit()
			}
		}()

		// handleError sets the error field on the given result. It returns
		// whether the caller should continue processing or break.
		handleError := func(result *Result, err error) bool {
			result.Error = err.Error()
			allResults = append(allResults, result)
			if tx {
				rollback = true // Will trigger the rollback.
				return false
			}
			return true
		}

		execer = db.sqlite3conn

		// Create the correct execution object, depending on whether a
		// transaction was requested.
		if tx {
			t, err = db.sqlite3conn.Begin()
			if err != nil {
				return err
			}
		}

		// Execute each query.
		for _, q := range queries {
			if q == "" {
				continue
			}

			result := &Result{}
			start := time.Now()

			r, err := execer.Exec(q, nil)
			if err != nil {
				if handleError(result, err) {
					continue
				}
				break
			}

			lid, err := r.LastInsertId()
			if err != nil {
				if handleError(result, err) {
					continue
				}
				break
			}
			result.LastInsertID = lid

			ra, err := r.RowsAffected()
			if err != nil {
				if handleError(result, err) {
					continue
				}
				break
			}
			result.RowsAffected = ra
			if xTime {
				result.Time = time.Now().Sub(start).Seconds()
			}
			allResults = append(allResults, result)
		}

		return nil
	}()

	return allResults, err
}

// Query executes queries that return rows, but don't modify the database.
func (db *DB) Query(queries []string, tx, xTime bool) ([]*Rows, error) {
	type Queryer interface {
		Query(query string, args []driver.Value) (driver.Rows, error)
	}

	var allRows []*Rows
	err := func() (err error) {
		var queryer Queryer
		var t driver.Tx
		defer func() {
			// XXX THIS DOESN'T ACTUALLY WORK! Might as WELL JUST COMMIT?
			if t != nil {
				if err != nil {
					t.Rollback()
					return
				}
				t.Commit()
			}
		}()

		queryer = db.sqlite3conn

		// Create the correct query object, depending on whether a
		// transaction was requested.
		if tx {
			t, err = db.sqlite3conn.Begin()
			if err != nil {
				return err
			}
		}

		for _, q := range queries {
			if q == "" {
				continue
			}

			rows := &Rows{}
			start := time.Now()

			rs, err := queryer.Query(q, nil)
			if err != nil {
				rows.Error = err.Error()
				allRows = append(allRows, rows)
				continue
			}
			defer rs.Close() // This adds to all defers, right? Nothing leaks? XXX Could consume memory. Perhaps anon would be best.
			columns := rs.Columns()

			rows.Columns = columns
			rows.Types = rs.(*sqlite3.SQLiteRows).DeclTypes()
			dest := make([]driver.Value, len(rows.Columns))
			for {
				err := rs.Next(dest)

				if err != nil {
					if err != io.EOF {
						rows.Error = err.Error()
					}
					break
				}

				values := make([]interface{}, len(rows.Columns))

				// Special case -- convert []uint8 to string. Perhaps this should be a config option.
				for i, v := range dest {
					switch u := v.(type) {
					case []uint8:
						values[i] = string(u)
					default:
						values[i] = u
					}
				}
				rows.Values = append(rows.Values, values)
			}
			if xTime {
				rows.Time = time.Now().Sub(start).Seconds()
			}
			allRows = append(allRows, rows)
		}

		return nil
	}()

	return allRows, err
}

// Backup writes a consistent snapshot of the database to the given file.
func (db *DB) Backup(path string) error {
	dstDB, err := Open(path)
	if err != nil {
		return err
	}

	bk, err := dstDB.sqlite3conn.Backup("main", db.sqlite3conn, "main")
	if err != nil {
		return err
	}

	for {
		done, err := bk.Step(-1)
		if err != nil {
			bk.Finish()
			return err
		}
		if done {
			break
		}
		time.Sleep(bkDelay * time.Millisecond)
	}

	if err := bk.Finish(); err != nil {
		return err
	}
	if err := dstDB.Close(); err != nil {
		return err
	}

	return nil
}
<|endoftext|>