<commit_before>// Copyright 2016 The Hugo Authors. All rights reserved.
//
// 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 commands

import (
	"fmt"
	"net"
	"net/http"
	"net/url"
	"os"
	"runtime"
	"strconv"
	"strings"
	"time"

	"mime"

	"github.com/spf13/afero"
	"github.com/spf13/cobra"
	"github.com/spf13/hugo/helpers"
	"github.com/spf13/hugo/hugofs"
	jww "github.com/spf13/jwalterweatherman"
	"github.com/spf13/viper"
)

var (
	disableLiveReload bool
	renderToDisk      bool
	serverAppend      bool
	serverInterface   string
	serverPort        int
	serverWatch       bool
)

//var serverCmdV *cobra.Command

var serverCmd = &cobra.Command{
	Use:     "server",
	Aliases: []string{"serve"},
	Short:   "A high performance webserver",
	Long: `Hugo provides its own webserver which builds and serves the site.
While hugo server is high performance, it is a webserver with limited options.
Many run it in production, but the standard behavior is for people to use it
in development and use a more full featured server such as Nginx or Caddy.

'hugo server' will avoid writing the rendered and served content to disk,
preferring to store it in memory.

By default hugo will also watch your files for any changes you make and
automatically rebuild the site. It will then live reload any open browser pages
and push the latest content to them. As most Hugo sites are built in a fraction
of a second, you will be able to save and see your changes nearly instantly.`,
	//RunE: server,
}

type filesOnlyFs struct {
	fs http.FileSystem
}

type noDirFile struct {
	http.File
}

func (fs filesOnlyFs) Open(name string) (http.File, error) {
	f, err := fs.fs.Open(name)
	if err != nil {
		return nil, err
	}
	return noDirFile{f}, nil
}

func (f noDirFile) Readdir(count int) ([]os.FileInfo, error) {
	return nil, nil
}

func init() {
	initHugoBuilderFlags(serverCmd)

	serverCmd.Flags().IntVarP(&serverPort, "port", "p", 1313, "port on which the server will listen")
	serverCmd.Flags().StringVarP(&serverInterface, "bind", "", "127.0.0.1", "interface to which the server will bind")
	serverCmd.Flags().BoolVarP(&serverWatch, "watch", "w", true, "watch filesystem for changes and recreate as needed")
	serverCmd.Flags().BoolVarP(&serverAppend, "appendPort", "", true, "append port to baseurl")
	serverCmd.Flags().BoolVar(&disableLiveReload, "disableLiveReload", false, "watch without enabling live browser reload on rebuild")
	serverCmd.Flags().BoolVar(&renderToDisk, "renderToDisk", false, "render to Destination path (default is render to memory & serve from there)")
	serverCmd.Flags().String("memstats", "", "log memory usage to this file")
	serverCmd.Flags().Int("meminterval", 100, "interval to poll memory usage (requires --memstats)")
	serverCmd.RunE = server

	mime.AddExtensionType(".json", "application/json; charset=utf-8")
	mime.AddExtensionType(".css", "text/css; charset=utf-8")

}

func server(cmd *cobra.Command, args []string) error {
	if err := InitializeConfig(serverCmd); err != nil {
		return err
	}

	if flagChanged(cmd.Flags(), "disableLiveReload") {
		viper.Set("DisableLiveReload", disableLiveReload)
	}

	if serverWatch {
		viper.Set("Watch", true)
	}

	if viper.GetBool("watch") {
		serverWatch = true
		watchConfig()
	}

	l, err := net.Listen("tcp", net.JoinHostPort(serverInterface, strconv.Itoa(serverPort)))
	if err == nil {
		l.Close()
	} else {
		if flagChanged(serverCmd.Flags(), "port") {
			// port set explicitly by user -- he/she probably meant it!
			return newSystemErrorF("Server startup failed: %s", err)
		}
		jww.ERROR.Println("port", serverPort, "already in use, attempting to use an available port")
		sp, err := helpers.FindAvailablePort()
		if err != nil {
			return newSystemError("Unable to find alternative port to use:", err)
		}
		serverPort = sp.Port
	}

	viper.Set("port", serverPort)

	BaseURL, err := fixURL(baseURL)
	if err != nil {
		return err
	}
	viper.Set("BaseURL", BaseURL)

	if err := memStats(); err != nil {
		jww.ERROR.Println("memstats error:", err)
	}

	// If a Destination is provided via flag write to disk
	if destination != "" {
		renderToDisk = true
	}

	// Hugo writes the output to memory instead of the disk
	if !renderToDisk {
		hugofs.SetDestination(new(afero.MemMapFs))
		// Rendering to memoryFS, publish to Root regardless of publishDir.
		viper.Set("PublishDir", "/")
	}

	if err := build(serverWatch); err != nil {
		return err
	}

	// Watch runs its own server as part of the routine
	if serverWatch {
		watchDirs := getDirList()
		baseWatchDir := viper.GetString("WorkingDir")
		for i, dir := range watchDirs {
			watchDirs[i], _ = helpers.GetRelativePath(dir, baseWatchDir)
		}

		rootWatchDirs := strings.Join(helpers.UniqueStrings(helpers.ExtractRootPaths(watchDirs)), ",")

		jww.FEEDBACK.Printf("Watching for changes in %s%s{%s}\n", baseWatchDir, helpers.FilePathSeparator, rootWatchDirs)
		err := NewWatcher(serverPort)

		if err != nil {
			return err
		}
	}

	serve(serverPort)

	return nil
}

func serve(port int) {
	if renderToDisk {
		jww.FEEDBACK.Println("Serving pages from " + helpers.AbsPathify(viper.GetString("PublishDir")))
	} else {
		jww.FEEDBACK.Println("Serving pages from memory")
	}

	httpFs := afero.NewHttpFs(hugofs.Destination())
	fs := filesOnlyFs{httpFs.Dir(helpers.AbsPathify(viper.GetString("PublishDir")))}
	fileserver := http.FileServer(fs)

	// We're only interested in the path
	u, err := url.Parse(viper.GetString("BaseURL"))
	if err != nil {
		jww.ERROR.Fatalf("Invalid BaseURL: %s", err)
	}
	if u.Path == "" || u.Path == "/" {
		http.Handle("/", fileserver)
	} else {
		http.Handle(u.Path, http.StripPrefix(u.Path, fileserver))
	}

	jww.FEEDBACK.Printf("Web Server is available at %s (bind address %s)\n", u.String(), serverInterface)
	fmt.Println("Press Ctrl+C to stop")

	endpoint := net.JoinHostPort(serverInterface, strconv.Itoa(port))
	err = http.ListenAndServe(endpoint, nil)
	if err != nil {
		jww.ERROR.Printf("Error: %s\n", err.Error())
		os.Exit(1)
	}
}

// fixURL massages the BaseURL into a form needed for serving
// all pages correctly.
func fixURL(s string) (string, error) {
	useLocalhost := false
	if s == "" {
		s = viper.GetString("BaseURL")
		useLocalhost = true
	}

	if !strings.HasSuffix(s, "/") {
		s = s + "/"
	}

	// do an initial parse of the input string
	u, err := url.Parse(s)
	if err != nil {
		return "", err
	}

	// if no Host is defined, then assume that no schema or double-slash were
	// present in the url.  Add a double-slash and make a best effort attempt.
	if u.Host == "" && s != "/" {
		s = "//" + s

		u, err = url.Parse(s)
		if err != nil {
			return "", err
		}
	}

	if useLocalhost {
		if u.Scheme == "https" {
			u.Scheme = "http"
		}
		u.Host = "localhost"
	}

	if serverAppend {
		if strings.Contains(u.Host, ":") {
			u.Host, _, err = net.SplitHostPort(u.Host)
			if err != nil {
				return "", fmt.Errorf("Failed to split BaseURL hostpost: %s", err)
			}
		}
		u.Host += fmt.Sprintf(":%d", serverPort)
	}

	return u.String(), nil
}

func memStats() error {
	memstats := serverCmd.Flags().Lookup("memstats").Value.String()
	if memstats != "" {
		interval, err := time.ParseDuration(serverCmd.Flags().Lookup("meminterval").Value.String())
		if err != nil {
			interval, _ = time.ParseDuration("100ms")
		}

		fileMemStats, err := os.Create(memstats)
		if err != nil {
			return err
		}

		fileMemStats.WriteString("# Time\tHeapSys\tHeapAlloc\tHeapIdle\tHeapReleased\n")

		go func() {
			var stats runtime.MemStats

			start := time.Now().UnixNano()

			for {
				runtime.ReadMemStats(&stats)
				if fileMemStats != nil {
					fileMemStats.WriteString(fmt.Sprintf("%d\t%d\t%d\t%d\t%d\n",
						(time.Now().UnixNano()-start)/1000000, stats.HeapSys, stats.HeapAlloc, stats.HeapIdle, stats.HeapReleased))
					time.Sleep(interval)
				} else {
					break
				}
			}
		}()
	}
	return nil
}
<commit_msg>Fix for meminterval not using specified interval<commit_after>// Copyright 2016 The Hugo Authors. All rights reserved.
//
// 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 commands

import (
	"fmt"
	"net"
	"net/http"
	"net/url"
	"os"
	"runtime"
	"strconv"
	"strings"
	"time"

	"mime"

	"github.com/spf13/afero"
	"github.com/spf13/cobra"
	"github.com/spf13/hugo/helpers"
	"github.com/spf13/hugo/hugofs"
	jww "github.com/spf13/jwalterweatherman"
	"github.com/spf13/viper"
)

var (
	disableLiveReload bool
	renderToDisk      bool
	serverAppend      bool
	serverInterface   string
	serverPort        int
	serverWatch       bool
)

//var serverCmdV *cobra.Command

var serverCmd = &cobra.Command{
	Use:     "server",
	Aliases: []string{"serve"},
	Short:   "A high performance webserver",
	Long: `Hugo provides its own webserver which builds and serves the site.
While hugo server is high performance, it is a webserver with limited options.
Many run it in production, but the standard behavior is for people to use it
in development and use a more full featured server such as Nginx or Caddy.

'hugo server' will avoid writing the rendered and served content to disk,
preferring to store it in memory.

By default hugo will also watch your files for any changes you make and
automatically rebuild the site. It will then live reload any open browser pages
and push the latest content to them. As most Hugo sites are built in a fraction
of a second, you will be able to save and see your changes nearly instantly.`,
	//RunE: server,
}

type filesOnlyFs struct {
	fs http.FileSystem
}

type noDirFile struct {
	http.File
}

func (fs filesOnlyFs) Open(name string) (http.File, error) {
	f, err := fs.fs.Open(name)
	if err != nil {
		return nil, err
	}
	return noDirFile{f}, nil
}

func (f noDirFile) Readdir(count int) ([]os.FileInfo, error) {
	return nil, nil
}

func init() {
	initHugoBuilderFlags(serverCmd)

	serverCmd.Flags().IntVarP(&serverPort, "port", "p", 1313, "port on which the server will listen")
	serverCmd.Flags().StringVarP(&serverInterface, "bind", "", "127.0.0.1", "interface to which the server will bind")
	serverCmd.Flags().BoolVarP(&serverWatch, "watch", "w", true, "watch filesystem for changes and recreate as needed")
	serverCmd.Flags().BoolVarP(&serverAppend, "appendPort", "", true, "append port to baseurl")
	serverCmd.Flags().BoolVar(&disableLiveReload, "disableLiveReload", false, "watch without enabling live browser reload on rebuild")
	serverCmd.Flags().BoolVar(&renderToDisk, "renderToDisk", false, "render to Destination path (default is render to memory & serve from there)")
	serverCmd.Flags().String("memstats", "", "log memory usage to this file")
	serverCmd.Flags().String("meminterval", "100ms", "interval to poll memory usage (requires --memstats), valid time units are \"ns\", \"us\" (or \"µs\"), \"ms\", \"s\", \"m\", \"h\".")

	serverCmd.RunE = server

	mime.AddExtensionType(".json", "application/json; charset=utf-8")
	mime.AddExtensionType(".css", "text/css; charset=utf-8")

}

func server(cmd *cobra.Command, args []string) error {
	if err := InitializeConfig(serverCmd); err != nil {
		return err
	}

	if flagChanged(cmd.Flags(), "disableLiveReload") {
		viper.Set("DisableLiveReload", disableLiveReload)
	}

	if serverWatch {
		viper.Set("Watch", true)
	}

	if viper.GetBool("watch") {
		serverWatch = true
		watchConfig()
	}

	l, err := net.Listen("tcp", net.JoinHostPort(serverInterface, strconv.Itoa(serverPort)))
	if err == nil {
		l.Close()
	} else {
		if flagChanged(serverCmd.Flags(), "port") {
			// port set explicitly by user -- he/she probably meant it!
			return newSystemErrorF("Server startup failed: %s", err)
		}
		jww.ERROR.Println("port", serverPort, "already in use, attempting to use an available port")
		sp, err := helpers.FindAvailablePort()
		if err != nil {
			return newSystemError("Unable to find alternative port to use:", err)
		}
		serverPort = sp.Port
	}

	viper.Set("port", serverPort)

	BaseURL, err := fixURL(baseURL)
	if err != nil {
		return err
	}
	viper.Set("BaseURL", BaseURL)

	if err := memStats(); err != nil {
		jww.ERROR.Println("memstats error:", err)
	}

	// If a Destination is provided via flag write to disk
	if destination != "" {
		renderToDisk = true
	}

	// Hugo writes the output to memory instead of the disk
	if !renderToDisk {
		hugofs.SetDestination(new(afero.MemMapFs))
		// Rendering to memoryFS, publish to Root regardless of publishDir.
		viper.Set("PublishDir", "/")
	}

	if err := build(serverWatch); err != nil {
		return err
	}

	// Watch runs its own server as part of the routine
	if serverWatch {
		watchDirs := getDirList()
		baseWatchDir := viper.GetString("WorkingDir")
		for i, dir := range watchDirs {
			watchDirs[i], _ = helpers.GetRelativePath(dir, baseWatchDir)
		}

		rootWatchDirs := strings.Join(helpers.UniqueStrings(helpers.ExtractRootPaths(watchDirs)), ",")

		jww.FEEDBACK.Printf("Watching for changes in %s%s{%s}\n", baseWatchDir, helpers.FilePathSeparator, rootWatchDirs)
		err := NewWatcher(serverPort)

		if err != nil {
			return err
		}
	}

	serve(serverPort)

	return nil
}

func serve(port int) {
	if renderToDisk {
		jww.FEEDBACK.Println("Serving pages from " + helpers.AbsPathify(viper.GetString("PublishDir")))
	} else {
		jww.FEEDBACK.Println("Serving pages from memory")
	}

	httpFs := afero.NewHttpFs(hugofs.Destination())
	fs := filesOnlyFs{httpFs.Dir(helpers.AbsPathify(viper.GetString("PublishDir")))}
	fileserver := http.FileServer(fs)

	// We're only interested in the path
	u, err := url.Parse(viper.GetString("BaseURL"))
	if err != nil {
		jww.ERROR.Fatalf("Invalid BaseURL: %s", err)
	}
	if u.Path == "" || u.Path == "/" {
		http.Handle("/", fileserver)
	} else {
		http.Handle(u.Path, http.StripPrefix(u.Path, fileserver))
	}

	jww.FEEDBACK.Printf("Web Server is available at %s (bind address %s)\n", u.String(), serverInterface)
	fmt.Println("Press Ctrl+C to stop")

	endpoint := net.JoinHostPort(serverInterface, strconv.Itoa(port))
	err = http.ListenAndServe(endpoint, nil)
	if err != nil {
		jww.ERROR.Printf("Error: %s\n", err.Error())
		os.Exit(1)
	}
}

// fixURL massages the BaseURL into a form needed for serving
// all pages correctly.
func fixURL(s string) (string, error) {
	useLocalhost := false
	if s == "" {
		s = viper.GetString("BaseURL")
		useLocalhost = true
	}

	if !strings.HasSuffix(s, "/") {
		s = s + "/"
	}

	// do an initial parse of the input string
	u, err := url.Parse(s)
	if err != nil {
		return "", err
	}

	// if no Host is defined, then assume that no schema or double-slash were
	// present in the url.  Add a double-slash and make a best effort attempt.
	if u.Host == "" && s != "/" {
		s = "//" + s

		u, err = url.Parse(s)
		if err != nil {
			return "", err
		}
	}

	if useLocalhost {
		if u.Scheme == "https" {
			u.Scheme = "http"
		}
		u.Host = "localhost"
	}

	if serverAppend {
		if strings.Contains(u.Host, ":") {
			u.Host, _, err = net.SplitHostPort(u.Host)
			if err != nil {
				return "", fmt.Errorf("Failed to split BaseURL hostpost: %s", err)
			}
		}
		u.Host += fmt.Sprintf(":%d", serverPort)
	}

	return u.String(), nil
}

func memStats() error {
	memstats := serverCmd.Flags().Lookup("memstats").Value.String()
	if memstats != "" {
		interval, err := time.ParseDuration(serverCmd.Flags().Lookup("meminterval").Value.String())
		if err != nil {
			interval, _ = time.ParseDuration("100ms")
		}

		fileMemStats, err := os.Create(memstats)
		if err != nil {
			return err
		}

		fileMemStats.WriteString("# Time\tHeapSys\tHeapAlloc\tHeapIdle\tHeapReleased\n")

		go func() {
			var stats runtime.MemStats

			start := time.Now().UnixNano()

			for {
				runtime.ReadMemStats(&stats)
				if fileMemStats != nil {
					fileMemStats.WriteString(fmt.Sprintf("%d\t%d\t%d\t%d\t%d\n",
						(time.Now().UnixNano()-start)/1000000, stats.HeapSys, stats.HeapAlloc, stats.HeapIdle, stats.HeapReleased))
					time.Sleep(interval)
				} else {
					break
				}
			}
		}()
	}
	return nil
}
<|endoftext|>


<commit_before>package driver_test

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"

	"github.com/pivotal-golang/lager/lagertest"

	"github.com/concourse/baggageclaim/fs"
	"github.com/concourse/baggageclaim/volume/driver"
)

var _ = Describe("BtrFS", func() {
	if runtime.GOOS != "linux" {
		fmt.Println("\x1b[33m*** skipping btrfs tests because non-linux ***\x1b[0m")
		return
	}

	var (
		tempDir    string
		volumeDir  string
		fsDriver   *driver.BtrFSDriver
		filesystem *fs.BtrfsFilesystem
	)

	BeforeEach(func() {
		var err error
		tempDir, err = ioutil.TempDir("", "baggageclaim_driver_test")
		Expect(err).NotTo(HaveOccurred())

		logger := lagertest.NewTestLogger("fs")

		imagePath := filepath.Join(tempDir, "image.img")
		volumeDir = filepath.Join(tempDir, "mountpoint")

		filesystem = fs.New(logger, imagePath, volumeDir)
		err = filesystem.Create(100 * 1024 * 1024)
		Expect(err).NotTo(HaveOccurred())

		fsDriver = driver.NewBtrFSDriver(logger, volumeDir)
	})

	AfterEach(func() {
		err := filesystem.Delete()
		Expect(err).NotTo(HaveOccurred())

		err = os.RemoveAll(tempDir)
		Expect(err).NotTo(HaveOccurred())
	})

	Describe("Lifecycle", func() {
		It("can create and delete a subvolume", func() {
			subvolumePath := filepath.Join(volumeDir, "subvolume")

			err := fsDriver.CreateVolume(subvolumePath)
			Expect(err).NotTo(HaveOccurred())

			Expect(subvolumePath).To(BeADirectory())

			checkSubvolume := exec.Command("btrfs", "subvolume", "show", subvolumePath)
			session, err := gexec.Start(checkSubvolume, GinkgoWriter, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())

			Eventually(session).Should(gbytes.Say(subvolumePath))
			Eventually(session).Should(gexec.Exit(0))

			err = fsDriver.DestroyVolume(subvolumePath)
			Expect(err).NotTo(HaveOccurred())

			Expect(subvolumePath).NotTo(BeADirectory())
		})

		It("can delete parent volume when it has subvolumes", func() {
			parentPath := filepath.Join(volumeDir, "parent")
			childPath := filepath.Join(parentPath, "volume", "child")
			grandchildPath := filepath.Join(childPath, "volume", "grandchild")

			err := os.MkdirAll(parentPath, os.ModePerm)
			Expect(err).NotTo(HaveOccurred())

			parentVolumePath := filepath.Join(parentPath, "volume")
			err = fsDriver.CreateVolume(parentVolumePath)
			Expect(err).NotTo(HaveOccurred())

			err = os.MkdirAll(filepath.Join(volumeDir, "sibling"), os.ModePerm)
			Expect(err).NotTo(HaveOccurred())

			siblingVolumePath := filepath.Join(volumeDir, "sibling", "volume")
			err = fsDriver.CreateVolume(siblingVolumePath)
			Expect(err).NotTo(HaveOccurred())

			err = os.MkdirAll(childPath, os.ModePerm)
			Expect(err).NotTo(HaveOccurred())
			childVolumePath := filepath.Join(childPath, "volume")

			err = fsDriver.CreateVolume(childVolumePath)
			Expect(err).NotTo(HaveOccurred())

			err = os.MkdirAll(grandchildPath, os.ModePerm)
			Expect(err).NotTo(HaveOccurred())
			grandchildVolumePath := filepath.Join(grandchildPath, "volume")
			err = fsDriver.CreateVolume(grandchildVolumePath)
			Expect(err).NotTo(HaveOccurred())

			err = fsDriver.DestroyVolume(parentVolumePath)
			Expect(err).NotTo(HaveOccurred())

			Expect(parentVolumePath).NotTo(BeADirectory())
			Expect(siblingVolumePath).To(BeADirectory())
		})
	})

	Describe("GetVolumeSize", func() {
		var parentVolumePath string
		var childVolumePath string

		BeforeEach(func() {
			parentVolumePath = filepath.Join(volumeDir, "parent-volume")
			err := fsDriver.CreateVolume(parentVolumePath)
			Expect(err).NotTo(HaveOccurred())

			bs := make([]byte, 4096)
			for i := 0; i < 4096; i++ {
				bs[i] = 'i'
			}
			err = ioutil.WriteFile(filepath.Join(parentVolumePath, "parent-stuff"), bs, os.ModePerm)
			Expect(err).NotTo(HaveOccurred())
		})

		AfterEach(func() {
			err := fsDriver.DestroyVolume(childVolumePath)
			Expect(err).NotTo(HaveOccurred())

			err = fsDriver.DestroyVolume(parentVolumePath)
			Expect(err).NotTo(HaveOccurred())
		})

		It("returns the approximate size of the volume at the given path", func() {
			childVolumePath = filepath.Join(volumeDir, "parent-volume", "child-volume")
			err := fsDriver.CreateVolume(childVolumePath)
			Expect(err).NotTo(HaveOccurred())

			size := 1024 * 1024 * 2
			bs := make([]byte, size) // 2 MiB
			for i := 0; i < size; i++ {
				bs[i] = 'i'
			}
			err = ioutil.WriteFile(filepath.Join(childVolumePath, "child-stuff"), bs, os.ModePerm)
			Expect(err).NotTo(HaveOccurred())

			Eventually(func() uint {
				GinkgoRecover()
				newSize, err := fsDriver.GetVolumeSize(childVolumePath)

				Expect(err).NotTo(HaveOccurred())
				return uint(newSize)
			}, 15*time.Second, 1*time.Second).Should(BeNumerically("~", size, float32(size)*.05))
		})
	})
})
<commit_msg>fix flaky btrfs size test<commit_after>package driver_test

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"

	"github.com/pivotal-golang/lager/lagertest"

	"github.com/concourse/baggageclaim/fs"
	"github.com/concourse/baggageclaim/volume/driver"
)

var _ = Describe("BtrFS", func() {
	if runtime.GOOS != "linux" {
		fmt.Println("\x1b[33m*** skipping btrfs tests because non-linux ***\x1b[0m")
		return
	}

	var (
		tempDir    string
		volumeDir  string
		fsDriver   *driver.BtrFSDriver
		filesystem *fs.BtrfsFilesystem
	)

	BeforeEach(func() {
		var err error
		tempDir, err = ioutil.TempDir("", "baggageclaim_driver_test")
		Expect(err).NotTo(HaveOccurred())

		logger := lagertest.NewTestLogger("fs")

		imagePath := filepath.Join(tempDir, "image.img")
		volumeDir = filepath.Join(tempDir, "mountpoint")

		filesystem = fs.New(logger, imagePath, volumeDir)
		err = filesystem.Create(100 * 1024 * 1024)
		Expect(err).NotTo(HaveOccurred())

		fsDriver = driver.NewBtrFSDriver(logger, volumeDir)
	})

	AfterEach(func() {
		err := filesystem.Delete()
		Expect(err).NotTo(HaveOccurred())

		err = os.RemoveAll(tempDir)
		Expect(err).NotTo(HaveOccurred())
	})

	Describe("Lifecycle", func() {
		It("can create and delete a subvolume", func() {
			subvolumePath := filepath.Join(volumeDir, "subvolume")

			err := fsDriver.CreateVolume(subvolumePath)
			Expect(err).NotTo(HaveOccurred())

			Expect(subvolumePath).To(BeADirectory())

			checkSubvolume := exec.Command("btrfs", "subvolume", "show", subvolumePath)
			session, err := gexec.Start(checkSubvolume, GinkgoWriter, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())

			Eventually(session).Should(gbytes.Say(subvolumePath))
			Eventually(session).Should(gexec.Exit(0))

			err = fsDriver.DestroyVolume(subvolumePath)
			Expect(err).NotTo(HaveOccurred())

			Expect(subvolumePath).NotTo(BeADirectory())
		})

		It("can delete parent volume when it has subvolumes", func() {
			parentPath := filepath.Join(volumeDir, "parent")
			childPath := filepath.Join(parentPath, "volume", "child")
			grandchildPath := filepath.Join(childPath, "volume", "grandchild")

			err := os.MkdirAll(parentPath, os.ModePerm)
			Expect(err).NotTo(HaveOccurred())

			parentVolumePath := filepath.Join(parentPath, "volume")
			err = fsDriver.CreateVolume(parentVolumePath)
			Expect(err).NotTo(HaveOccurred())

			err = os.MkdirAll(filepath.Join(volumeDir, "sibling"), os.ModePerm)
			Expect(err).NotTo(HaveOccurred())

			siblingVolumePath := filepath.Join(volumeDir, "sibling", "volume")
			err = fsDriver.CreateVolume(siblingVolumePath)
			Expect(err).NotTo(HaveOccurred())

			err = os.MkdirAll(childPath, os.ModePerm)
			Expect(err).NotTo(HaveOccurred())
			childVolumePath := filepath.Join(childPath, "volume")

			err = fsDriver.CreateVolume(childVolumePath)
			Expect(err).NotTo(HaveOccurred())

			err = os.MkdirAll(grandchildPath, os.ModePerm)
			Expect(err).NotTo(HaveOccurred())
			grandchildVolumePath := filepath.Join(grandchildPath, "volume")
			err = fsDriver.CreateVolume(grandchildVolumePath)
			Expect(err).NotTo(HaveOccurred())

			err = fsDriver.DestroyVolume(parentVolumePath)
			Expect(err).NotTo(HaveOccurred())

			Expect(parentVolumePath).NotTo(BeADirectory())
			Expect(siblingVolumePath).To(BeADirectory())
		})
	})

	Describe("GetVolumeSize", func() {
		var parentVolumePath string
		var childVolumePath string

		BeforeEach(func() {
			parentVolumePath = filepath.Join(volumeDir, "parent-volume")
			err := fsDriver.CreateVolume(parentVolumePath)
			Expect(err).NotTo(HaveOccurred())

			bs := make([]byte, 4096)
			for i := 0; i < 4096; i++ {
				bs[i] = 'i'
			}
			err = ioutil.WriteFile(filepath.Join(parentVolumePath, "parent-stuff"), bs, os.ModePerm)
			Expect(err).NotTo(HaveOccurred())
		})

		AfterEach(func() {
			err := fsDriver.DestroyVolume(childVolumePath)
			Expect(err).NotTo(HaveOccurred())

			err = fsDriver.DestroyVolume(parentVolumePath)
			Expect(err).NotTo(HaveOccurred())
		})

		It("returns the size of the volume at the given path", func() {
			childVolumePath = filepath.Join(volumeDir, "parent-volume", "child-volume")
			err := fsDriver.CreateVolume(childVolumePath)
			Expect(err).NotTo(HaveOccurred())

			originalSize, err := fsDriver.GetVolumeSize(childVolumePath)
			Expect(err).NotTo(HaveOccurred())

			size := 1024 * 1024 * 2
			bs := make([]byte, size) // 2 MiB
			for i := 0; i < size; i++ {
				bs[i] = 'i'
			}
			err = ioutil.WriteFile(filepath.Join(childVolumePath, "child-stuff"), bs, os.ModePerm)
			Expect(err).NotTo(HaveOccurred())

			timeout := 1 * time.Minute // btrfs periodic commit happens every 30 seconds
			Eventually(func() uint {
				GinkgoRecover()
				newSize, err := fsDriver.GetVolumeSize(childVolumePath)

				Expect(err).NotTo(HaveOccurred())
				return newSize
			}, timeout, 1*time.Second).Should(Equal(uint(size) + originalSize))
		})
	})
})
<|endoftext|>


<commit_before>package driver_test

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"

	"github.com/pivotal-golang/lager/lagertest"

	"github.com/concourse/baggageclaim/fs"
	"github.com/concourse/baggageclaim/volume/driver"
)

var _ = Describe("BtrFS", func() {
	if runtime.GOOS != "linux" {
		fmt.Println("\x1b[33m*** skipping btrfs tests because non-linux ***\x1b[0m")
		return
	}

	var (
		tempDir    string
		volumeDir  string
		fsDriver   *driver.BtrFSDriver
		filesystem *fs.BtrfsFilesystem
	)

	BeforeEach(func() {
		var err error
		tempDir, err = ioutil.TempDir("", "baggageclaim_driver_test")
		Expect(err).NotTo(HaveOccurred())

		logger := lagertest.NewTestLogger("fs")

		imagePath := filepath.Join(tempDir, "image.img")
		volumeDir = filepath.Join(tempDir, "mountpoint")

		filesystem = fs.New(logger, imagePath, volumeDir)
		err = filesystem.Create(100 * 1024 * 1024)
		Expect(err).NotTo(HaveOccurred())

		fsDriver = driver.NewBtrFSDriver(logger, volumeDir)
	})

	AfterEach(func() {
		err := filesystem.Delete()
		Expect(err).NotTo(HaveOccurred())

		err = os.RemoveAll(tempDir)
		Expect(err).NotTo(HaveOccurred())
	})

	Describe("Lifecycle", func() {
		It("can create and delete a subvolume", func() {
			subvolumePath := filepath.Join(volumeDir, "subvolume")

			err := fsDriver.CreateVolume(subvolumePath)
			Expect(err).NotTo(HaveOccurred())

			Expect(subvolumePath).To(BeADirectory())

			checkSubvolume := exec.Command("btrfs", "subvolume", "show", subvolumePath)
			session, err := gexec.Start(checkSubvolume, GinkgoWriter, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())

			Eventually(session).Should(gbytes.Say(subvolumePath))
			Eventually(session).Should(gexec.Exit(0))

			err = fsDriver.DestroyVolume(subvolumePath)
			Expect(err).NotTo(HaveOccurred())

			Expect(subvolumePath).NotTo(BeADirectory())
		})

		It("can delete parent volume when it has subvolumes", func() {
			parentPath := filepath.Join(volumeDir, "parent")
			childPath := filepath.Join(parentPath, "volume", "child")
			grandchildPath := filepath.Join(childPath, "volume", "grandchild")

			err := os.MkdirAll(parentPath, os.ModePerm)
			Expect(err).NotTo(HaveOccurred())

			parentVolumePath := filepath.Join(parentPath, "volume")
			err = fsDriver.CreateVolume(parentVolumePath)
			Expect(err).NotTo(HaveOccurred())

			err = os.MkdirAll(filepath.Join(volumeDir, "sibling"), os.ModePerm)
			Expect(err).NotTo(HaveOccurred())

			siblingVolumePath := filepath.Join(volumeDir, "sibling", "volume")
			err = fsDriver.CreateVolume(siblingVolumePath)
			Expect(err).NotTo(HaveOccurred())

			err = os.MkdirAll(childPath, os.ModePerm)
			Expect(err).NotTo(HaveOccurred())
			childVolumePath := filepath.Join(childPath, "volume")

			err = fsDriver.CreateVolume(childVolumePath)
			Expect(err).NotTo(HaveOccurred())

			err = os.MkdirAll(grandchildPath, os.ModePerm)
			Expect(err).NotTo(HaveOccurred())
			grandchildVolumePath := filepath.Join(grandchildPath, "volume")
			err = fsDriver.CreateVolume(grandchildVolumePath)
			Expect(err).NotTo(HaveOccurred())

			err = fsDriver.DestroyVolume(parentVolumePath)
			Expect(err).NotTo(HaveOccurred())

			Expect(parentVolumePath).NotTo(BeADirectory())
			Expect(siblingVolumePath).To(BeADirectory())
		})
	})

	Describe("GetVolumeSize", func() {
		var parentVolumePath string
		var childVolumePath string

		BeforeEach(func() {
			parentVolumePath = filepath.Join(volumeDir, "parent-volume")
			err := fsDriver.CreateVolume(parentVolumePath)
			Expect(err).NotTo(HaveOccurred())

			bs := make([]byte, 4096)
			for i := 0; i < 4096; i++ {
				bs[i] = 'i'
			}
			err = ioutil.WriteFile(filepath.Join(parentVolumePath, "parent-stuff"), bs, os.ModePerm)
			Expect(err).NotTo(HaveOccurred())
		})

		AfterEach(func() {
			err := fsDriver.DestroyVolume(childVolumePath)
			Expect(err).NotTo(HaveOccurred())

			err = fsDriver.DestroyVolume(parentVolumePath)
			Expect(err).NotTo(HaveOccurred())
		})

		It("returns the size of the volume at the given path", func() {
			childVolumePath = filepath.Join(volumeDir, "parent-volume", "child-volume")
			err := fsDriver.CreateVolume(childVolumePath)
			Expect(err).NotTo(HaveOccurred())

			originalSize, err := fsDriver.GetVolumeSize(childVolumePath)
			Expect(err).NotTo(HaveOccurred())

			size := 1024 * 1024 * 2
			bs := make([]byte, size) // 2 MiB
			for i := 0; i < size; i++ {
				bs[i] = 'i'
			}
			err = ioutil.WriteFile(filepath.Join(childVolumePath, "child-stuff"), bs, os.ModePerm)
			Expect(err).NotTo(HaveOccurred())

			timeout := 1 * time.Minute // btrfs periodic commit happens every 30 seconds
			Eventually(func() uint {
				GinkgoRecover()
				newSize, err := fsDriver.GetVolumeSize(childVolumePath)

				Expect(err).NotTo(HaveOccurred())
				return newSize
			}, timeout, 1*time.Second).Should(Equal(uint(size) + originalSize))
		})
	})
})
<commit_msg>bump btrfs volume size test timeout<commit_after>package driver_test

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"

	"github.com/pivotal-golang/lager/lagertest"

	"github.com/concourse/baggageclaim/fs"
	"github.com/concourse/baggageclaim/volume/driver"
)

var _ = Describe("BtrFS", func() {
	if runtime.GOOS != "linux" {
		fmt.Println("\x1b[33m*** skipping btrfs tests because non-linux ***\x1b[0m")
		return
	}

	var (
		tempDir    string
		volumeDir  string
		fsDriver   *driver.BtrFSDriver
		filesystem *fs.BtrfsFilesystem
	)

	BeforeEach(func() {
		var err error
		tempDir, err = ioutil.TempDir("", "baggageclaim_driver_test")
		Expect(err).NotTo(HaveOccurred())

		logger := lagertest.NewTestLogger("fs")

		imagePath := filepath.Join(tempDir, "image.img")
		volumeDir = filepath.Join(tempDir, "mountpoint")

		filesystem = fs.New(logger, imagePath, volumeDir)
		err = filesystem.Create(100 * 1024 * 1024)
		Expect(err).NotTo(HaveOccurred())

		fsDriver = driver.NewBtrFSDriver(logger, volumeDir)
	})

	AfterEach(func() {
		err := filesystem.Delete()
		Expect(err).NotTo(HaveOccurred())

		err = os.RemoveAll(tempDir)
		Expect(err).NotTo(HaveOccurred())
	})

	Describe("Lifecycle", func() {
		It("can create and delete a subvolume", func() {
			subvolumePath := filepath.Join(volumeDir, "subvolume")

			err := fsDriver.CreateVolume(subvolumePath)
			Expect(err).NotTo(HaveOccurred())

			Expect(subvolumePath).To(BeADirectory())

			checkSubvolume := exec.Command("btrfs", "subvolume", "show", subvolumePath)
			session, err := gexec.Start(checkSubvolume, GinkgoWriter, GinkgoWriter)
			Expect(err).NotTo(HaveOccurred())

			Eventually(session).Should(gbytes.Say(subvolumePath))
			Eventually(session).Should(gexec.Exit(0))

			err = fsDriver.DestroyVolume(subvolumePath)
			Expect(err).NotTo(HaveOccurred())

			Expect(subvolumePath).NotTo(BeADirectory())
		})

		It("can delete parent volume when it has subvolumes", func() {
			parentPath := filepath.Join(volumeDir, "parent")
			childPath := filepath.Join(parentPath, "volume", "child")
			grandchildPath := filepath.Join(childPath, "volume", "grandchild")

			err := os.MkdirAll(parentPath, os.ModePerm)
			Expect(err).NotTo(HaveOccurred())

			parentVolumePath := filepath.Join(parentPath, "volume")
			err = fsDriver.CreateVolume(parentVolumePath)
			Expect(err).NotTo(HaveOccurred())

			err = os.MkdirAll(filepath.Join(volumeDir, "sibling"), os.ModePerm)
			Expect(err).NotTo(HaveOccurred())

			siblingVolumePath := filepath.Join(volumeDir, "sibling", "volume")
			err = fsDriver.CreateVolume(siblingVolumePath)
			Expect(err).NotTo(HaveOccurred())

			err = os.MkdirAll(childPath, os.ModePerm)
			Expect(err).NotTo(HaveOccurred())
			childVolumePath := filepath.Join(childPath, "volume")

			err = fsDriver.CreateVolume(childVolumePath)
			Expect(err).NotTo(HaveOccurred())

			err = os.MkdirAll(grandchildPath, os.ModePerm)
			Expect(err).NotTo(HaveOccurred())
			grandchildVolumePath := filepath.Join(grandchildPath, "volume")
			err = fsDriver.CreateVolume(grandchildVolumePath)
			Expect(err).NotTo(HaveOccurred())

			err = fsDriver.DestroyVolume(parentVolumePath)
			Expect(err).NotTo(HaveOccurred())

			Expect(parentVolumePath).NotTo(BeADirectory())
			Expect(siblingVolumePath).To(BeADirectory())
		})
	})

	Describe("GetVolumeSize", func() {
		var parentVolumePath string
		var childVolumePath string

		BeforeEach(func() {
			parentVolumePath = filepath.Join(volumeDir, "parent-volume")
			err := fsDriver.CreateVolume(parentVolumePath)
			Expect(err).NotTo(HaveOccurred())

			bs := make([]byte, 4096)
			for i := 0; i < 4096; i++ {
				bs[i] = 'i'
			}
			err = ioutil.WriteFile(filepath.Join(parentVolumePath, "parent-stuff"), bs, os.ModePerm)
			Expect(err).NotTo(HaveOccurred())
		})

		AfterEach(func() {
			err := fsDriver.DestroyVolume(childVolumePath)
			Expect(err).NotTo(HaveOccurred())

			err = fsDriver.DestroyVolume(parentVolumePath)
			Expect(err).NotTo(HaveOccurred())
		})

		It("returns the size of the volume at the given path", func() {
			childVolumePath = filepath.Join(volumeDir, "parent-volume", "child-volume")
			err := fsDriver.CreateVolume(childVolumePath)
			Expect(err).NotTo(HaveOccurred())

			originalSize, err := fsDriver.GetVolumeSize(childVolumePath)
			Expect(err).NotTo(HaveOccurred())

			size := 1024 * 1024 * 2
			bs := make([]byte, size) // 2 MiB
			for i := 0; i < size; i++ {
				bs[i] = 'i'
			}
			err = ioutil.WriteFile(filepath.Join(childVolumePath, "child-stuff"), bs, os.ModePerm)
			Expect(err).NotTo(HaveOccurred())

			timeout := 2 * time.Minute // btrfs periodic commit happens every 30 seconds
			Eventually(func() uint {
				GinkgoRecover()
				newSize, err := fsDriver.GetVolumeSize(childVolumePath)

				Expect(err).NotTo(HaveOccurred())
				return newSize
			}, timeout, 1*time.Second).Should(Equal(uint(size) + originalSize))
		})
	})
})
<|endoftext|>


<commit_before>// Copyright 2016 The StudyGolang Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// http://studygolang.com
// Author:polaris	polaris@studygolang.com

package logic

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/polaris1119/config"
	"github.com/polaris1119/goutils"
	"github.com/polaris1119/logger"
	"github.com/polaris1119/times"
)

const (
	WsMsgNotify = iota // 通知消息
	WsMsgOnline        // 发送在线用户数（和需要时也发历史最高）
)

const MessageQueueLen = 5

type Message struct {
	Type int         `json:"type"`
	Body interface{} `json:"body"`
}

func NewMessage(msgType int, msgBody interface{}) *Message {
	return &Message{
		Type: msgType,
		Body: msgBody,
	}
}

type UserData struct {
	// 该用户收到的消息（key为serverId）
	serverMsgQueue map[int]chan *Message
	lastAccessTime time.Time
	onlineDuartion time.Duration

	rwMutex sync.RWMutex
}

func (this *UserData) Len() int {
	this.rwMutex.RLock()
	defer this.rwMutex.RUnlock()
	return len(this.serverMsgQueue)
}

func (this *UserData) MessageQueue(serverId int) chan *Message {
	this.rwMutex.RLock()
	defer this.rwMutex.RUnlock()
	return this.serverMsgQueue[serverId]
}

func (this *UserData) Remove(serverId int) {
	this.rwMutex.Lock()
	defer this.rwMutex.Unlock()
	delete(this.serverMsgQueue, serverId)
}

func (this *UserData) InitMessageQueue(serverId int) {
	this.rwMutex.Lock()
	defer this.rwMutex.Unlock()
	this.serverMsgQueue[serverId] = make(chan *Message, MessageQueueLen)
}

func (this *UserData) SendMessage(message *Message) {
	this.rwMutex.RLock()
	defer this.rwMutex.RUnlock()

	for serverId, messageQueue := range this.serverMsgQueue {
		// 有可能用户已经退出，导致 messageQueue 满，阻塞
		if len(messageQueue) < MessageQueueLen {
			messageQueue <- message
		} else {
			logger.Infoln("server_id:", serverId, "had close")

			delete(this.serverMsgQueue, serverId)
		}
	}
}

// 用于 expvar 统计信息
type LoginUser struct {
	Uid            int    `json:"uid"`
	LastAccessTime string `json:"last_access_time"`
	OnlineDuartion string `json:"online_duration"`
}
type LoginUserSlice []*LoginUser

func (self LoginUserSlice) String() string {
	b, err := json.Marshal(self)
	if err != nil {
		return "[]"
	}

	return string(b)
}

var Book = &book{users: make(map[int]*UserData), uids: make(map[int]struct{})}

type book struct {
	users map[int]*UserData
	// 登录用户
	uids    map[int]struct{}
	rwMutex sync.RWMutex
}

// 增加一个用户到book中（有可能是用户的另一个请求）
// user为UID或IP地址的int表示
func (this *book) AddUser(user, serverId int, isUid bool) *UserData {
	var userData *UserData
	var ok bool
	this.rwMutex.Lock()
	if userData, ok = this.users[user]; ok {
		this.rwMutex.Unlock()

		userData.InitMessageQueue(serverId)
		userData.onlineDuartion += time.Now().Sub(userData.lastAccessTime)
		userData.lastAccessTime = time.Now()
	} else {
		userData = &UserData{
			serverMsgQueue: map[int]chan *Message{serverId: make(chan *Message, MessageQueueLen)},
			lastAccessTime: time.Now(),
		}
		this.users[user] = userData
		if isUid {
			this.uids[user] = struct{}{}
		}
		length := len(this.users)

		this.rwMutex.Unlock()

		onlineInfo := map[string]int{"online": length}
		// 在线人数超过历史最高
		if length > MaxOnlineNum() {
			maxRwMu.Lock()
			maxOnlineNum = length
			onlineInfo["maxonline"] = maxOnlineNum
			maxRwMu.Unlock()
			saveMaxOnlineNum()
		}
		// 广播给其他人：有新用户进来，包括可能的新历史最高
		message := NewMessage(WsMsgOnline, onlineInfo)
		go this.BroadcastToOthersMessage(message, user)
	}

	return userData
}

// 删除用户
func (this *book) DelUser(user, serverId int, isUid bool) {
	this.rwMutex.Lock()
	defer this.rwMutex.Unlock()

	// 自己只有一个页面建立websocket连接
	if this.users[user].Len() == 1 {
		delete(this.users, user)
		if isUid {
			delete(this.uids, user)
		}
	} else {
		this.users[user].Remove(serverId)
	}
}

// 判断用户是否还在线（user 有可能是IP）
func (this *book) UserIsOnline(user int) bool {
	this.rwMutex.RLock()
	defer this.rwMutex.RUnlock()
	if _, ok := this.users[user]; ok {
		return true
	}
	return false
}

// 判断注册用户是否还在线（user 有可能是IP）
func (this *book) RegUserIsOnline(uid int) bool {
	this.rwMutex.RLock()
	defer this.rwMutex.RUnlock()
	if _, ok := this.uids[uid]; ok {
		return true
	}
	return false
}

// 在线用户数
func (this *book) Len() int {
	this.rwMutex.RLock()
	defer this.rwMutex.RUnlock()
	return len(this.users)
}

// 在线注册会员数
func (this *book) LoginLen() int {
	this.rwMutex.RLock()
	defer this.rwMutex.RUnlock()
	return len(this.uids)
}

func (this *book) LoginUserData() LoginUserSlice {
	this.rwMutex.RLock()
	defer this.rwMutex.RUnlock()

	loginUserData := LoginUserSlice(make([]*LoginUser, 0, len(this.uids)))
	for uid := range this.uids {
		user := this.users[uid]
		loginUserData = append(loginUserData, &LoginUser{
			Uid:            uid,
			LastAccessTime: times.Format("Y-m-d H:i:s", user.lastAccessTime),
			OnlineDuartion: user.onlineDuartion.String(),
		})
	}

	return loginUserData
}

// 给某个用户发送一条消息
func (this *book) PostMessage(uid int, message *Message) {
	this.rwMutex.RLock()
	defer this.rwMutex.RUnlock()
	if userData, ok := this.users[uid]; ok {
		logger.Infoln("post message to", uid, message)
		go userData.SendMessage(message)
	}
}

// 给所有用户广播消息
func (this *book) BroadcastAllUsersMessage(message *Message) {
	logger.Infoln("BroadcastAllUsersMessage message", message)

	this.rwMutex.RLock()
	defer this.rwMutex.RUnlock()
	for _, userData := range this.users {
		userData.SendMessage(message)
	}
}

// 给除了自己的其他用户广播消息
func (this *book) BroadcastToOthersMessage(message *Message, myself int) {
	logger.Infoln("BroadcastToOthersMessage message", message)

	this.rwMutex.RLock()
	defer this.rwMutex.RUnlock()
	for uid, userData := range this.users {
		if uid == myself {
			continue
		}
		userData.SendMessage(message)
	}
}

var (
	// 保存历史最大在线用户数
	maxOnlineNum int
	maxRwMu      sync.RWMutex
)

func initMaxOnlineNum() {
	maxRwMu.Lock()
	defer maxRwMu.Unlock()
	if maxOnlineNum == 0 {
		data, err := ioutil.ReadFile(getDataFile())
		if err != nil {
			logger.Errorln("read data file error:", err)
			return
		}
		maxOnlineNum = goutils.MustInt(strings.TrimSpace(string(data)))
	}
}

// 获得历史最高在线人数
func MaxOnlineNum() int {
	initMaxOnlineNum()
	maxRwMu.RLock()
	defer maxRwMu.RUnlock()
	return maxOnlineNum
}

func saveMaxOnlineNum() {
	data := []byte(strconv.Itoa(MaxOnlineNum()))
	err := ioutil.WriteFile(getDataFile(), data, 0777)
	if err != nil {
		logger.Errorln("write data file error:", err)
		return
	}
}

var dataFile string

func getDataFile() string {
	if dataFile != "" {
		return dataFile
	}
	dataFile = config.ConfigFile.MustValue("global", "data_path")
	if !filepath.IsAbs(dataFile) {
		dataFile = config.ROOT + "/" + dataFile
	}
	// 文件夹不存在，则创建
	dataPath := filepath.Dir(dataFile)
	if err := os.MkdirAll(dataPath, 0777); err != nil {
		logger.Errorln("MkdirAll error:", err)
	}
	return dataFile
}
<commit_msg>用户进入记录log<commit_after>// Copyright 2016 The StudyGolang Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// http://studygolang.com
// Author:polaris	polaris@studygolang.com

package logic

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/polaris1119/config"
	"github.com/polaris1119/goutils"
	"github.com/polaris1119/logger"
	"github.com/polaris1119/times"
)

const (
	WsMsgNotify = iota // 通知消息
	WsMsgOnline        // 发送在线用户数（和需要时也发历史最高）
)

const MessageQueueLen = 5

type Message struct {
	Type int         `json:"type"`
	Body interface{} `json:"body"`
}

func NewMessage(msgType int, msgBody interface{}) *Message {
	return &Message{
		Type: msgType,
		Body: msgBody,
	}
}

type UserData struct {
	// 该用户收到的消息（key为serverId）
	serverMsgQueue map[int]chan *Message
	lastAccessTime time.Time
	onlineDuartion time.Duration

	rwMutex sync.RWMutex
}

func (this *UserData) Len() int {
	this.rwMutex.RLock()
	defer this.rwMutex.RUnlock()
	return len(this.serverMsgQueue)
}

func (this *UserData) MessageQueue(serverId int) chan *Message {
	this.rwMutex.RLock()
	defer this.rwMutex.RUnlock()
	return this.serverMsgQueue[serverId]
}

func (this *UserData) Remove(serverId int) {
	this.rwMutex.Lock()
	defer this.rwMutex.Unlock()
	delete(this.serverMsgQueue, serverId)
}

func (this *UserData) InitMessageQueue(serverId int) {
	this.rwMutex.Lock()
	defer this.rwMutex.Unlock()
	this.serverMsgQueue[serverId] = make(chan *Message, MessageQueueLen)
}

func (this *UserData) SendMessage(message *Message) {
	this.rwMutex.RLock()
	defer this.rwMutex.RUnlock()

	for serverId, messageQueue := range this.serverMsgQueue {
		// 有可能用户已经退出，导致 messageQueue 满，阻塞
		if len(messageQueue) < MessageQueueLen {
			messageQueue <- message
		} else {
			logger.Infoln("server_id:", serverId, "had close")

			delete(this.serverMsgQueue, serverId)
		}
	}
}

// 用于 expvar 统计信息
type LoginUser struct {
	Uid            int    `json:"uid"`
	LastAccessTime string `json:"last_access_time"`
	OnlineDuartion string `json:"online_duration"`
}
type LoginUserSlice []*LoginUser

func (self LoginUserSlice) String() string {
	b, err := json.Marshal(self)
	if err != nil {
		return "[]"
	}

	return string(b)
}

var Book = &book{users: make(map[int]*UserData), uids: make(map[int]struct{})}

type book struct {
	users map[int]*UserData
	// 登录用户
	uids    map[int]struct{}
	rwMutex sync.RWMutex
}

// 增加一个用户到book中（有可能是用户的另一个请求）
// user为UID或IP地址的int表示
func (this *book) AddUser(user, serverId int, isUid bool) *UserData {
	var userData *UserData
	var ok bool
	this.rwMutex.Lock()
	if userData, ok = this.users[user]; ok {
		this.rwMutex.Unlock()

		logger.Infoln("user:", user, "had enter")

		userData.InitMessageQueue(serverId)
		userData.onlineDuartion += time.Now().Sub(userData.lastAccessTime)
		userData.lastAccessTime = time.Now()
	} else {
		userData = &UserData{
			serverMsgQueue: map[int]chan *Message{serverId: make(chan *Message, MessageQueueLen)},
			lastAccessTime: time.Now(),
		}
		this.users[user] = userData
		if isUid {
			this.uids[user] = struct{}{}
		}
		length := len(this.users)

		this.rwMutex.Unlock()

		onlineInfo := map[string]int{"online": length}
		// 在线人数超过历史最高
		if length > MaxOnlineNum() {
			maxRwMu.Lock()
			maxOnlineNum = length
			onlineInfo["maxonline"] = maxOnlineNum
			maxRwMu.Unlock()
			saveMaxOnlineNum()
		}
		// 广播给其他人：有新用户进来，包括可能的新历史最高
		message := NewMessage(WsMsgOnline, onlineInfo)
		go this.BroadcastToOthersMessage(message, user)
	}

	return userData
}

// 删除用户
func (this *book) DelUser(user, serverId int, isUid bool) {
	this.rwMutex.Lock()
	defer this.rwMutex.Unlock()

	// 自己只有一个页面建立websocket连接
	if this.users[user].Len() == 1 {
		delete(this.users, user)
		if isUid {
			delete(this.uids, user)
		}
	} else {
		this.users[user].Remove(serverId)
	}
}

// 判断用户是否还在线（user 有可能是IP）
func (this *book) UserIsOnline(user int) bool {
	this.rwMutex.RLock()
	defer this.rwMutex.RUnlock()
	if _, ok := this.users[user]; ok {
		return true
	}
	return false
}

// 判断注册用户是否还在线（user 有可能是IP）
func (this *book) RegUserIsOnline(uid int) bool {
	this.rwMutex.RLock()
	defer this.rwMutex.RUnlock()
	if _, ok := this.uids[uid]; ok {
		return true
	}
	return false
}

// 在线用户数
func (this *book) Len() int {
	this.rwMutex.RLock()
	defer this.rwMutex.RUnlock()
	return len(this.users)
}

// 在线注册会员数
func (this *book) LoginLen() int {
	this.rwMutex.RLock()
	defer this.rwMutex.RUnlock()
	return len(this.uids)
}

func (this *book) LoginUserData() LoginUserSlice {
	this.rwMutex.RLock()
	defer this.rwMutex.RUnlock()

	loginUserData := LoginUserSlice(make([]*LoginUser, 0, len(this.uids)))
	for uid := range this.uids {
		user := this.users[uid]
		loginUserData = append(loginUserData, &LoginUser{
			Uid:            uid,
			LastAccessTime: times.Format("Y-m-d H:i:s", user.lastAccessTime),
			OnlineDuartion: user.onlineDuartion.String(),
		})
	}

	return loginUserData
}

// 给某个用户发送一条消息
func (this *book) PostMessage(uid int, message *Message) {
	this.rwMutex.RLock()
	defer this.rwMutex.RUnlock()
	if userData, ok := this.users[uid]; ok {
		logger.Infoln("post message to", uid, message)
		go userData.SendMessage(message)
	}
}

// 给所有用户广播消息
func (this *book) BroadcastAllUsersMessage(message *Message) {
	logger.Infoln("BroadcastAllUsersMessage message", message)

	this.rwMutex.RLock()
	defer this.rwMutex.RUnlock()
	for _, userData := range this.users {
		userData.SendMessage(message)
	}
}

// 给除了自己的其他用户广播消息
func (this *book) BroadcastToOthersMessage(message *Message, myself int) {
	logger.Infoln("BroadcastToOthersMessage message", message)

	this.rwMutex.RLock()
	defer this.rwMutex.RUnlock()
	for uid, userData := range this.users {
		if uid == myself {
			continue
		}
		userData.SendMessage(message)
	}
}

var (
	// 保存历史最大在线用户数
	maxOnlineNum int
	maxRwMu      sync.RWMutex
)

func initMaxOnlineNum() {
	maxRwMu.Lock()
	defer maxRwMu.Unlock()
	if maxOnlineNum == 0 {
		data, err := ioutil.ReadFile(getDataFile())
		if err != nil {
			logger.Errorln("read data file error:", err)
			return
		}
		maxOnlineNum = goutils.MustInt(strings.TrimSpace(string(data)))
	}
}

// 获得历史最高在线人数
func MaxOnlineNum() int {
	initMaxOnlineNum()
	maxRwMu.RLock()
	defer maxRwMu.RUnlock()
	return maxOnlineNum
}

func saveMaxOnlineNum() {
	data := []byte(strconv.Itoa(MaxOnlineNum()))
	err := ioutil.WriteFile(getDataFile(), data, 0777)
	if err != nil {
		logger.Errorln("write data file error:", err)
		return
	}
}

var dataFile string

func getDataFile() string {
	if dataFile != "" {
		return dataFile
	}
	dataFile = config.ConfigFile.MustValue("global", "data_path")
	if !filepath.IsAbs(dataFile) {
		dataFile = config.ROOT + "/" + dataFile
	}
	// 文件夹不存在，则创建
	dataPath := filepath.Dir(dataFile)
	if err := os.MkdirAll(dataPath, 0777); err != nil {
		logger.Errorln("MkdirAll error:", err)
	}
	return dataFile
}
<|endoftext|>


<commit_before>/*
Package compose provides a Go wrapper around Docker Compose, useful for integration testing.

	var composeYML =`
	test_mockserver:
	  container_name: ms
	  image: jamesdbloom/mockserver
	  ports:
	    - "10000:1080"
	    - "1090"
	test_postgres:
	  container_name: pg
	  image: postgres
	  ports:
	    - "5432"

	compose, err := compose.Start(composeYML, true, true)
	if err != nil {
		panic(err)
	}
	defer compose.Kill()
*/
package compose

import (
	"fmt"
	"log"
	"os"
	"regexp"
	"strings"
)

// Compose is the main type exported by the package, used to interact with a running Docker Compose configuration.
type Compose struct {
	FileName   string
	Containers map[string]*Container
}

var (
	logger           = log.New(os.Stdout, "go-compose: ", log.LstdFlags)
	replaceEnvRegexp = regexp.MustCompile("\\$\\{[^\\}]+\\}")
	composeUpRegexp  = regexp.MustCompile("(?m:^docker start <- \\(u'(.*)'\\)$)")
)

const (
	composeProjectName = "compose"
)

// Start starts a Docker Compose configuration.
// If forcePull is true, it attempts do pull newer versions of the images.
// If rmFirst is true, it attempts to kill and delete containers before starting new ones.
func Start(dockerComposeYML string, forcePull, rmFirst bool) (*Compose, error) {
	logger.Println("initializing...")
	dockerComposeYML = replaceEnv(dockerComposeYML)

	fName, err := writeTmp(dockerComposeYML)
	if err != nil {
		return nil, err
	}

	ids, err := startCompose(fName, forcePull, rmFirst)
	if err != nil {
		return nil, err
	}

	containers := make(map[string]*Container)

	for _, id := range ids {
		container, err := Inspect(id)
		if err != nil {
			return nil, err
		}
		if !container.State.Running {
			return nil, fmt.Errorf("compose: container '%v' is not running", container.Name)
		}
		containers[container.Name[1:]] = container
	}

	return &Compose{FileName: fName, Containers: containers}, nil
}

// MustStart is like Start, but panics on error.
func MustStart(dockerComposeYML string, forcePull, killFirst bool) *Compose {
	compose, err := Start(dockerComposeYML, forcePull, killFirst)
	if err != nil {
		panic(err)
	}
	return compose
}

// Kill kills any running containers for the current configuration.
func (c *Compose) Kill() error {
	logger.Println("killing containers...")
	_, err := runCompose(c.FileName, "kill")
	if err == nil {
		logger.Println("containers killed")
		return nil
	}
	return fmt.Errorf("compose: error killing containers: %v", err)
}

// MustKill is like Kill, but panics on error.
func (c *Compose) MustKill() {
	if err := c.Kill(); err != nil {
		panic(err)
	}
}

func replaceEnv(dockerComposeYML string) string {
	return replaceEnvRegexp.ReplaceAllStringFunc(dockerComposeYML, replaceEnvFunc)
}

func replaceEnvFunc(s string) string {
	return os.Getenv(strings.TrimSpace(s[2 : len(s)-1]))
}

func startCompose(fName string, forcePull, rmFirst bool) ([]string, error) {
	if forcePull {
		logger.Println("pulling images...")
		if _, err := runCompose(fName, "pull"); err != nil {
			return nil, fmt.Errorf("compose: error pulling images: %v", err)
		}
	}

	if rmFirst {
		logger.Println("removing stale containers...")
		_, err := runCompose(fName, "kill")
		if err != nil {
			return nil, fmt.Errorf("compose: error killing stale containers: %v", err)
		}
		logger.Println("removing stale containers...")
		_, err = runCompose(fName, "rm", "--force")
		if err != nil {
			return nil, fmt.Errorf("compose: error removing stale containers: %v", err)
		}
	}

	logger.Println("starting containers...")
	out, err := runCompose(fName, "--verbose", "up", "-d")
	fmt.Println(out)
	if err != nil {
		return nil, fmt.Errorf("compose: error starting containers: %v", err)
	}
	logger.Println("containers started")

	matches := composeUpRegexp.FindAllStringSubmatch(out, -1)
	ids := make([]string, 0, len(matches))
	for _, match := range matches {
		ids = append(ids, match[1])
	}

	return ids, nil
}

func runCompose(fName string, otherArgs ...string) (string, error) {
	args := []string{"-f", fName, "-p", composeProjectName}
	args = append(args, otherArgs...)
	return runCmd("docker-compose", args...)
}
<commit_msg>Remove unnecessary command output.<commit_after>/*
Package compose provides a Go wrapper around Docker Compose, useful for integration testing.

	var composeYML =`
	test_mockserver:
	  container_name: ms
	  image: jamesdbloom/mockserver
	  ports:
	    - "10000:1080"
	    - "1090"
	test_postgres:
	  container_name: pg
	  image: postgres
	  ports:
	    - "5432"

	compose, err := compose.Start(composeYML, true, true)
	if err != nil {
		panic(err)
	}
	defer compose.Kill()
*/
package compose

import (
	"fmt"
	"log"
	"os"
	"regexp"
	"strings"
)

// Compose is the main type exported by the package, used to interact with a running Docker Compose configuration.
type Compose struct {
	FileName   string
	Containers map[string]*Container
}

var (
	logger           = log.New(os.Stdout, "go-compose: ", log.LstdFlags)
	replaceEnvRegexp = regexp.MustCompile("\\$\\{[^\\}]+\\}")
	composeUpRegexp  = regexp.MustCompile("(?m:^docker start <- \\(u'(.*)'\\)$)")
)

const (
	composeProjectName = "compose"
)

// Start starts a Docker Compose configuration.
// If forcePull is true, it attempts do pull newer versions of the images.
// If rmFirst is true, it attempts to kill and delete containers before starting new ones.
func Start(dockerComposeYML string, forcePull, rmFirst bool) (*Compose, error) {
	logger.Println("initializing...")
	dockerComposeYML = replaceEnv(dockerComposeYML)

	fName, err := writeTmp(dockerComposeYML)
	if err != nil {
		return nil, err
	}

	ids, err := startCompose(fName, forcePull, rmFirst)
	if err != nil {
		return nil, err
	}

	containers := make(map[string]*Container)

	for _, id := range ids {
		container, err := Inspect(id)
		if err != nil {
			return nil, err
		}
		if !container.State.Running {
			return nil, fmt.Errorf("compose: container '%v' is not running", container.Name)
		}
		containers[container.Name[1:]] = container
	}

	return &Compose{FileName: fName, Containers: containers}, nil
}

// MustStart is like Start, but panics on error.
func MustStart(dockerComposeYML string, forcePull, killFirst bool) *Compose {
	compose, err := Start(dockerComposeYML, forcePull, killFirst)
	if err != nil {
		panic(err)
	}
	return compose
}

// Kill kills any running containers for the current configuration.
func (c *Compose) Kill() error {
	logger.Println("killing containers...")
	_, err := runCompose(c.FileName, "kill")
	if err == nil {
		logger.Println("containers killed")
		return nil
	}
	return fmt.Errorf("compose: error killing containers: %v", err)
}

// MustKill is like Kill, but panics on error.
func (c *Compose) MustKill() {
	if err := c.Kill(); err != nil {
		panic(err)
	}
}

func replaceEnv(dockerComposeYML string) string {
	return replaceEnvRegexp.ReplaceAllStringFunc(dockerComposeYML, replaceEnvFunc)
}

func replaceEnvFunc(s string) string {
	return os.Getenv(strings.TrimSpace(s[2 : len(s)-1]))
}

func startCompose(fName string, forcePull, rmFirst bool) ([]string, error) {
	if forcePull {
		logger.Println("pulling images...")
		if _, err := runCompose(fName, "pull"); err != nil {
			return nil, fmt.Errorf("compose: error pulling images: %v", err)
		}
	}

	if rmFirst {
		logger.Println("removing stale containers...")
		_, err := runCompose(fName, "kill")
		if err != nil {
			return nil, fmt.Errorf("compose: error killing stale containers: %v", err)
		}
		logger.Println("removing stale containers...")
		_, err = runCompose(fName, "rm", "--force")
		if err != nil {
			return nil, fmt.Errorf("compose: error removing stale containers: %v", err)
		}
	}

	logger.Println("starting containers...")
	out, err := runCompose(fName, "--verbose", "up", "-d")
	if err != nil {
		return nil, fmt.Errorf("compose: error starting containers: %v", err)
	}
	logger.Println("containers started")

	matches := composeUpRegexp.FindAllStringSubmatch(out, -1)
	ids := make([]string, 0, len(matches))
	for _, match := range matches {
		ids = append(ids, match[1])
	}

	return ids, nil
}

func runCompose(fName string, otherArgs ...string) (string, error) {
	args := []string{"-f", fName, "-p", composeProjectName}
	args = append(args, otherArgs...)
	return runCmd("docker-compose", args...)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"log"
	"os"
	"runtime"
	"time"

	"github.com/jasonmoo/wc"
	"github.com/reducedb/bloom"
	"github.com/reducedb/bloom/scalable"
)

var (
	intersection = flag.Bool("i", false, "calculate the intersection")
	diff         = flag.Bool("d", false, "calculate the difference")
	union        = flag.Bool("u", false, "calculate the union")

	count = flag.Bool("c", false, "output counts of each token on non-large unions")

	// bloom processing
	large           = flag.Bool("large", false, "use bloom filters for large data size (may be lossy)")
	estimated_lines = flag.Uint64("estimated_lines", 0, "estimate used to size bloom filters (set this to avoid prescan)")

	// options
	trim          = flag.Bool("trim", false, "trim each line")
	match_regex   = flag.String("match", "", "only process matching lines")
	capture_regex = flag.String("capture", "", "only process captured data")

	// for fs opts
	devnull     = flag.Bool("devnull", false, "do not output tokens, just counts")
	buffer_size = flag.Int("buffer_size", 1<<20, "buffered io chunk size")

	// totals
	total_tokens_emitted  uint64
	total_bytes_processed uint64
	total_lines_scanned   uint64
	total_lines_matched   uint64
)

func init() {
	runtime.GOMAXPROCS(runtime.NumCPU())
}

func main() {

	flag.Parse()

	if !*intersection && !*diff && !*union {
		fmt.Println(`Usage: tt -[i,d,u] [-c] [-trim] [-match "regex"] [-capture "regex"] [-large [-estimated_lines N]] file1 file2[ file3..]`)
		flag.PrintDefaults()
		os.Exit(1)
	}

	start := time.Now()

	var stdout WriteFlusher

	if *devnull {
		stdout = new(DevNullWriter)
	} else {
		// buffered io
		stdout = bufio.NewWriterSize(os.Stdout, *buffer_size)
	}

	defer func() {
		stdout.Flush()
		fmt.Fprintln(os.Stderr, "** Token Report **")
		fmt.Fprintln(os.Stderr, "Lines scanned: ", total_lines_scanned)
		if *match_regex != "" {
			fmt.Fprintln(os.Stderr, "Lines matched: ", total_lines_matched)
		}
		fmt.Fprintln(os.Stderr, "Tokens emitted: ", total_tokens_emitted)
		fmt.Fprintln(os.Stderr, "Time: ", time.Since(start))
	}()

	file_paths := flag.Args()

	fmt.Fprintln(os.Stderr, "tt starting up")

	// if no estimate supplied, count lines
	if *large && *estimated_lines == 0 {

		var bytes_to_process uint64

		for _, file_path := range file_paths {

			file, err := os.Open(file_path)
			if err != nil {
				log.Fatal(err)
			}

			counter := wc.NewCounter(file)
			err = counter.Count(false, true, true, false)
			if err != nil {
				log.Fatal(err)
			}

			*estimated_lines += counter.Lines
			bytes_to_process += counter.Bytes

			file.Close()
		}

		fmt.Fprintln(os.Stderr, "Bytes to process: ", bytes_to_process)
		fmt.Fprintln(os.Stderr, "Lines to process: ", *estimated_lines)
	}

	if *large {

		if *union {

			unique_set := NewScalableBloom(*estimated_lines)

			for _, file_path := range file_paths {

				e, err := NewEmitter(file_path, *match_regex, *capture_regex, *buffer_size)
				if err != nil {
					log.Fatal(err)
				}

				for e.Scan() {
					token := e.Bytes()
					if !unique_set.Check(token) {
						total_tokens_emitted++
						stdout.Write(token)
						stdout.WriteByte('\n')
						unique_set.Add(token)
					}
				}

				e.Close()

				total_lines_scanned += e.LinesScanned

			}

			return
		}

		// multi file handling below
		sets := make([]bloom.Bloom, len(file_paths))

		// may require throttling due to disk thrashing
		// initial scan to fill the bloom filters
		for i, file_path := range file_paths {

			set := NewScalableBloom(*estimated_lines)

			e, err := NewEmitter(file_path, *match_regex, *capture_regex, *buffer_size)
			if err != nil {
				log.Fatal(err)
			}

			for e.Scan() {
				set.Add(e.Bytes())
			}

			e.Close()

			sets[i] = set

		}

		// do the work
		switch {

		// unique set of tokens that exist in all files
		case *intersection:

			echoed_set := NewScalableBloom(*estimated_lines)

			for _, file_path := range file_paths {

				e, err := NewEmitter(file_path, *match_regex, *capture_regex, *buffer_size)
				if err != nil {
					log.Fatal(err)
				}

			NEXT_TOKEN:
				for e.Scan() {

					token := e.Bytes()

					if echoed_set.Check(token) {
						goto NEXT_TOKEN
					}

					for _, set := range sets {
						if !set.Check(token) {
							goto NEXT_TOKEN
						}
					}

					total_tokens_emitted++
					stdout.Write(token)
					stdout.WriteByte('\n')
					echoed_set.Add(token)

				}

				total_lines_scanned += e.LinesScanned

				e.Close()

			}

		// unique set of tokens not in the intersection
		case *diff:

			echoed_set := NewScalableBloom(*estimated_lines)

			for _, file_path := range file_paths {

				e, err := NewEmitter(file_path, *match_regex, *capture_regex, *buffer_size)
				if err != nil {
					log.Fatal(err)
				}

				for e.Scan() {

					token := e.Bytes()

					if echoed_set.Check(token) {
						continue
					}

					for _, set := range sets {
						if !set.Check(token) {
							total_tokens_emitted++
							stdout.Write(token)
							stdout.WriteByte('\n')
							echoed_set.Add(token)
						}
					}

				}

				total_lines_scanned += e.LinesScanned

				e.Close()

			}
		}

		// defaults to map solution
	} else {

		if *union {

			unique_set := make(map[string]int)

			for _, file_path := range file_paths {

				e, err := NewEmitter(file_path, *match_regex, *capture_regex, *buffer_size)
				if err != nil {
					log.Fatal(err)
				}

				for e.Scan() {
					unique_set[e.Text()]++
				}

				total_lines_scanned += e.LinesScanned

				e.Close()

			}

			if *count {
				for token, ct := range unique_set {
					total_tokens_emitted++
					fmt.Fprintf(stdout, "%d: %s\n", ct, token)
				}
			} else {
				for token, _ := range unique_set {
					total_tokens_emitted++
					stdout.WriteString(token)
					stdout.WriteByte('\n')
				}
			}

			return
		}

		// multi file handling below
		sets := make([]map[string]bool, len(file_paths))

		// may require throttling due to disk thrashing
		// initial scan to fill the bloom filters
		for i, file_path := range file_paths {

			set := make(map[string]bool)

			e, err := NewEmitter(file_path, *match_regex, *capture_regex, *buffer_size)
			if err != nil {
				log.Fatal(err)
			}

			for e.Scan() {
				set[e.Text()] = true
			}

			e.Close()

			sets[i] = set

		}

		// do the work
		switch {

		// unique set of tokens that exist in all files
		case *intersection:

			echoed_set := make(map[string]bool)

			for _, file_path := range file_paths {

				e, err := NewEmitter(file_path, *match_regex, *capture_regex, *buffer_size)
				if err != nil {
					log.Fatal(err)
				}

			NEXT_TOKEN2:
				for e.Scan() {

					token := e.Text()

					if _, echoed := echoed_set[token]; echoed {
						goto NEXT_TOKEN2
					}

					for _, set := range sets {
						if _, in_this_set := set[token]; !in_this_set {
							goto NEXT_TOKEN2
						}
					}

					total_tokens_emitted++
					stdout.WriteString(token)
					stdout.WriteByte('\n')

					echoed_set[token] = true

				}

				total_lines_scanned += e.LinesScanned

				e.Close()

			}

		// unique set of tokens not in the intersection
		case *diff:

			echoed_set := make(map[string]bool)

			for _, file_path := range file_paths {

				e, err := NewEmitter(file_path, *match_regex, *capture_regex, *buffer_size)
				if err != nil {
					log.Fatal(err)
				}

				for e.Scan() {

					token := e.Text()

					if _, echoed := echoed_set[token]; echoed {
						continue
					}

					for _, set := range sets {
						if _, in_this_set := set[token]; !in_this_set {
							total_tokens_emitted++
							stdout.WriteString(token)
							stdout.WriteByte('\n')

							echoed_set[token] = true
							break
						}
					}

				}

				total_lines_scanned += e.LinesScanned

				e.Close()

			}
		}

	}

}

func NewScalableBloom(size uint64) bloom.Bloom {
	return scalable.New(uint(size))
}
<commit_msg>fixing deps<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"log"
	"os"
	"runtime"
	"time"

	"github.com/dataence/bloom"
	"github.com/dataence/bloom/scalable"
	"github.com/jasonmoo/wc"
)

var (
	intersection = flag.Bool("i", false, "calculate the intersection")
	diff         = flag.Bool("d", false, "calculate the difference")
	union        = flag.Bool("u", false, "calculate the union")

	count = flag.Bool("c", false, "output counts of each token on non-large unions")

	// bloom processing
	large           = flag.Bool("large", false, "use bloom filters for large data size (may be lossy)")
	estimated_lines = flag.Uint64("estimated_lines", 0, "estimate used to size bloom filters (set this to avoid prescan)")

	// options
	trim          = flag.Bool("trim", false, "trim each line")
	match_regex   = flag.String("match", "", "only process matching lines")
	capture_regex = flag.String("capture", "", "only process captured data")

	// for fs opts
	devnull     = flag.Bool("devnull", false, "do not output tokens, just counts")
	buffer_size = flag.Int("buffer_size", 1<<20, "buffered io chunk size")

	// totals
	total_tokens_emitted  uint64
	total_bytes_processed uint64
	total_lines_scanned   uint64
	total_lines_matched   uint64
)

func init() {
	runtime.GOMAXPROCS(runtime.NumCPU())
}

func main() {

	flag.Parse()

	if !*intersection && !*diff && !*union {
		fmt.Println(`Usage: tt -[i,d,u] [-c] [-trim] [-match "regex"] [-capture "regex"] [-large [-estimated_lines N]] file1 file2[ file3..]`)
		flag.PrintDefaults()
		os.Exit(1)
	}

	start := time.Now()

	var stdout WriteFlusher

	if *devnull {
		stdout = new(DevNullWriter)
	} else {
		// buffered io
		stdout = bufio.NewWriterSize(os.Stdout, *buffer_size)
	}

	defer func() {
		stdout.Flush()
		fmt.Fprintln(os.Stderr, "** Token Report **")
		fmt.Fprintln(os.Stderr, "Lines scanned: ", total_lines_scanned)
		if *match_regex != "" {
			fmt.Fprintln(os.Stderr, "Lines matched: ", total_lines_matched)
		}
		fmt.Fprintln(os.Stderr, "Tokens emitted: ", total_tokens_emitted)
		fmt.Fprintln(os.Stderr, "Time: ", time.Since(start))
	}()

	file_paths := flag.Args()

	fmt.Fprintln(os.Stderr, "tt starting up")

	// if no estimate supplied, count lines
	if *large && *estimated_lines == 0 {

		var bytes_to_process uint64

		for _, file_path := range file_paths {

			file, err := os.Open(file_path)
			if err != nil {
				log.Fatal(err)
			}

			counter := wc.NewCounter(file)
			err = counter.Count(false, true, true, false)
			if err != nil {
				log.Fatal(err)
			}

			*estimated_lines += counter.Lines
			bytes_to_process += counter.Bytes

			file.Close()
		}

		fmt.Fprintln(os.Stderr, "Bytes to process: ", bytes_to_process)
		fmt.Fprintln(os.Stderr, "Lines to process: ", *estimated_lines)
	}

	if *large {

		if *union {

			unique_set := NewScalableBloom(*estimated_lines)

			for _, file_path := range file_paths {

				e, err := NewEmitter(file_path, *match_regex, *capture_regex, *buffer_size)
				if err != nil {
					log.Fatal(err)
				}

				for e.Scan() {
					token := e.Bytes()
					if !unique_set.Check(token) {
						total_tokens_emitted++
						stdout.Write(token)
						stdout.WriteByte('\n')
						unique_set.Add(token)
					}
				}

				e.Close()

				total_lines_scanned += e.LinesScanned

			}

			return
		}

		// multi file handling below
		sets := make([]bloom.Bloom, len(file_paths))

		// may require throttling due to disk thrashing
		// initial scan to fill the bloom filters
		for i, file_path := range file_paths {

			set := NewScalableBloom(*estimated_lines)

			e, err := NewEmitter(file_path, *match_regex, *capture_regex, *buffer_size)
			if err != nil {
				log.Fatal(err)
			}

			for e.Scan() {
				set.Add(e.Bytes())
			}

			e.Close()

			sets[i] = set

		}

		// do the work
		switch {

		// unique set of tokens that exist in all files
		case *intersection:

			echoed_set := NewScalableBloom(*estimated_lines)

			for _, file_path := range file_paths {

				e, err := NewEmitter(file_path, *match_regex, *capture_regex, *buffer_size)
				if err != nil {
					log.Fatal(err)
				}

			NEXT_TOKEN:
				for e.Scan() {

					token := e.Bytes()

					if echoed_set.Check(token) {
						goto NEXT_TOKEN
					}

					for _, set := range sets {
						if !set.Check(token) {
							goto NEXT_TOKEN
						}
					}

					total_tokens_emitted++
					stdout.Write(token)
					stdout.WriteByte('\n')
					echoed_set.Add(token)

				}

				total_lines_scanned += e.LinesScanned

				e.Close()

			}

		// unique set of tokens not in the intersection
		case *diff:

			echoed_set := NewScalableBloom(*estimated_lines)

			for _, file_path := range file_paths {

				e, err := NewEmitter(file_path, *match_regex, *capture_regex, *buffer_size)
				if err != nil {
					log.Fatal(err)
				}

				for e.Scan() {

					token := e.Bytes()

					if echoed_set.Check(token) {
						continue
					}

					for _, set := range sets {
						if !set.Check(token) {
							total_tokens_emitted++
							stdout.Write(token)
							stdout.WriteByte('\n')
							echoed_set.Add(token)
						}
					}

				}

				total_lines_scanned += e.LinesScanned

				e.Close()

			}
		}

		// defaults to map solution
	} else {

		if *union {

			unique_set := make(map[string]int)

			for _, file_path := range file_paths {

				e, err := NewEmitter(file_path, *match_regex, *capture_regex, *buffer_size)
				if err != nil {
					log.Fatal(err)
				}

				for e.Scan() {
					unique_set[e.Text()]++
				}

				total_lines_scanned += e.LinesScanned

				e.Close()

			}

			if *count {
				for token, ct := range unique_set {
					total_tokens_emitted++
					fmt.Fprintf(stdout, "%d: %s\n", ct, token)
				}
			} else {
				for token, _ := range unique_set {
					total_tokens_emitted++
					stdout.WriteString(token)
					stdout.WriteByte('\n')
				}
			}

			return
		}

		// multi file handling below
		sets := make([]map[string]bool, len(file_paths))

		// may require throttling due to disk thrashing
		// initial scan to fill the bloom filters
		for i, file_path := range file_paths {

			set := make(map[string]bool)

			e, err := NewEmitter(file_path, *match_regex, *capture_regex, *buffer_size)
			if err != nil {
				log.Fatal(err)
			}

			for e.Scan() {
				set[e.Text()] = true
			}

			e.Close()

			sets[i] = set

		}

		// do the work
		switch {

		// unique set of tokens that exist in all files
		case *intersection:

			echoed_set := make(map[string]bool)

			for _, file_path := range file_paths {

				e, err := NewEmitter(file_path, *match_regex, *capture_regex, *buffer_size)
				if err != nil {
					log.Fatal(err)
				}

			NEXT_TOKEN2:
				for e.Scan() {

					token := e.Text()

					if _, echoed := echoed_set[token]; echoed {
						goto NEXT_TOKEN2
					}

					for _, set := range sets {
						if _, in_this_set := set[token]; !in_this_set {
							goto NEXT_TOKEN2
						}
					}

					total_tokens_emitted++
					stdout.WriteString(token)
					stdout.WriteByte('\n')

					echoed_set[token] = true

				}

				total_lines_scanned += e.LinesScanned

				e.Close()

			}

		// unique set of tokens not in the intersection
		case *diff:

			echoed_set := make(map[string]bool)

			for _, file_path := range file_paths {

				e, err := NewEmitter(file_path, *match_regex, *capture_regex, *buffer_size)
				if err != nil {
					log.Fatal(err)
				}

				for e.Scan() {

					token := e.Text()

					if _, echoed := echoed_set[token]; echoed {
						continue
					}

					for _, set := range sets {
						if _, in_this_set := set[token]; !in_this_set {
							total_tokens_emitted++
							stdout.WriteString(token)
							stdout.WriteByte('\n')

							echoed_set[token] = true
							break
						}
					}

				}

				total_lines_scanned += e.LinesScanned

				e.Close()

			}
		}

	}

}

func NewScalableBloom(size uint64) bloom.Bloom {
	return scalable.New(uint(size))
}
<|endoftext|>


<commit_before>package gost

import (
	"crypto/tls"
	"github.com/golang/glog"
	"github.com/gorilla/websocket"
	"net"
	"net/http"
	"net/http/httputil"
	"time"
)

type WebsocketServer struct {
	Addr     string
	Base     *ProxyServer
	Handler  http.Handler
	upgrader websocket.Upgrader
}

func NewWebsocketServer(base *ProxyServer) *WebsocketServer {
	return &WebsocketServer{
		Addr: base.Node.Addr,
		Base: base,
		upgrader: websocket.Upgrader{
			ReadBufferSize:    1024,
			WriteBufferSize:   1024,
			CheckOrigin:       func(r *http.Request) bool { return true },
			EnableCompression: true,
		},
	}
}

// Default websocket server handler
func (s *WebsocketServer) HandleRequest(w http.ResponseWriter, r *http.Request) {
	glog.V(LINFO).Infof("[ws] %s - %s", r.RemoteAddr, s.Addr)
	if glog.V(LDEBUG) {
		dump, _ := httputil.DumpRequest(r, false)
		glog.V(LDEBUG).Infof("[ws] %s - %s\n%s", r.RemoteAddr, s.Addr, string(dump))
	}
	conn, err := s.upgrader.Upgrade(w, r, nil)
	if err != nil {
		glog.V(LERROR).Infof("[ws] %s - %s : %s", r.RemoteAddr, s.Addr, err)
		return
	}
	s.Base.handleConn(WebsocketServerConn(conn))
}

func (s *WebsocketServer) ListenAndServe() error {
	mux := http.NewServeMux()
	if s.Handler == nil {
		s.Handler = http.HandlerFunc(s.HandleRequest)
	}
	mux.Handle("/ws", s.Handler)
	return http.ListenAndServe(s.Addr, mux)
}

func (s *WebsocketServer) ListenAndServeTLS(config *tls.Config) error {
	mux := http.NewServeMux()
	if s.Handler == nil {
		s.Handler = http.HandlerFunc(s.HandleRequest)
	}
	mux.Handle("/ws", s.Handler)
	server := &http.Server{
		Addr:      s.Addr,
		Handler:   mux,
		TLSConfig: config,
	}
	return server.ListenAndServeTLS("", "")
}

type WebsocketConn struct {
	conn *websocket.Conn
	rb   []byte
}

func WebsocketClientConn(url string, conn net.Conn, config *tls.Config) (*WebsocketConn, error) {
	dialer := websocket.Dialer{
		ReadBufferSize:    1024,
		WriteBufferSize:   1024,
		TLSClientConfig:   config,
		HandshakeTimeout:  DialTimeout,
		EnableCompression: true,
		NetDial: func(net, addr string) (net.Conn, error) {
			return conn, nil
		},
	}

	c, resp, err := dialer.Dial(url, nil)
	if err != nil {
		return nil, err
	}
	resp.Body.Close()
	return &WebsocketConn{conn: c}, nil
}

func WebsocketServerConn(conn *websocket.Conn) *WebsocketConn {
	conn.EnableWriteCompression(true)
	return &WebsocketConn{
		conn: conn,
	}
}

func (c *WebsocketConn) Read(b []byte) (n int, err error) {
	if len(c.rb) == 0 {
		_, c.rb, err = c.conn.ReadMessage()
	}
	n = copy(b, c.rb)
	c.rb = c.rb[n:]
	return
}

func (c *WebsocketConn) Write(b []byte) (n int, err error) {
	err = c.conn.WriteMessage(websocket.BinaryMessage, b)
	n = len(b)
	return
}

func (c *WebsocketConn) Close() error {
	return c.conn.Close()
}

func (c *WebsocketConn) LocalAddr() net.Addr {
	return c.conn.LocalAddr()
}

func (c *WebsocketConn) RemoteAddr() net.Addr {
	return c.conn.RemoteAddr()
}

func (conn *WebsocketConn) SetDeadline(t time.Time) error {
	if err := conn.SetReadDeadline(t); err != nil {
		return err
	}
	return conn.SetWriteDeadline(t)
}
func (c *WebsocketConn) SetReadDeadline(t time.Time) error {
	return c.conn.SetReadDeadline(t)
}

func (c *WebsocketConn) SetWriteDeadline(t time.Time) error {
	return c.conn.SetWriteDeadline(t)
}
<commit_msg>update websocket package<commit_after>package gost

import (
	"crypto/tls"
	"github.com/golang/glog"
	"gopkg.in/gorilla/websocket.v1"
	"net"
	"net/http"
	"net/http/httputil"
	"time"
)

type WebsocketServer struct {
	Addr     string
	Base     *ProxyServer
	Handler  http.Handler
	upgrader websocket.Upgrader
}

func NewWebsocketServer(base *ProxyServer) *WebsocketServer {
	return &WebsocketServer{
		Addr: base.Node.Addr,
		Base: base,
		upgrader: websocket.Upgrader{
			ReadBufferSize:    1024,
			WriteBufferSize:   1024,
			CheckOrigin:       func(r *http.Request) bool { return true },
			EnableCompression: true,
		},
	}
}

// Default websocket server handler
func (s *WebsocketServer) HandleRequest(w http.ResponseWriter, r *http.Request) {
	glog.V(LINFO).Infof("[ws] %s - %s", r.RemoteAddr, s.Addr)
	if glog.V(LDEBUG) {
		dump, _ := httputil.DumpRequest(r, false)
		glog.V(LDEBUG).Infof("[ws] %s - %s\n%s", r.RemoteAddr, s.Addr, string(dump))
	}
	conn, err := s.upgrader.Upgrade(w, r, nil)
	if err != nil {
		glog.V(LERROR).Infof("[ws] %s - %s : %s", r.RemoteAddr, s.Addr, err)
		return
	}
	s.Base.handleConn(WebsocketServerConn(conn))
}

func (s *WebsocketServer) ListenAndServe() error {
	mux := http.NewServeMux()
	if s.Handler == nil {
		s.Handler = http.HandlerFunc(s.HandleRequest)
	}
	mux.Handle("/ws", s.Handler)
	return http.ListenAndServe(s.Addr, mux)
}

func (s *WebsocketServer) ListenAndServeTLS(config *tls.Config) error {
	mux := http.NewServeMux()
	if s.Handler == nil {
		s.Handler = http.HandlerFunc(s.HandleRequest)
	}
	mux.Handle("/ws", s.Handler)
	server := &http.Server{
		Addr:      s.Addr,
		Handler:   mux,
		TLSConfig: config,
	}
	return server.ListenAndServeTLS("", "")
}

type WebsocketConn struct {
	conn *websocket.Conn
	rb   []byte
}

func WebsocketClientConn(url string, conn net.Conn, config *tls.Config) (*WebsocketConn, error) {
	dialer := websocket.Dialer{
		ReadBufferSize:    1024,
		WriteBufferSize:   1024,
		TLSClientConfig:   config,
		HandshakeTimeout:  DialTimeout,
		EnableCompression: true,
		NetDial: func(net, addr string) (net.Conn, error) {
			return conn, nil
		},
	}

	c, resp, err := dialer.Dial(url, nil)
	if err != nil {
		return nil, err
	}
	resp.Body.Close()
	return &WebsocketConn{conn: c}, nil
}

func WebsocketServerConn(conn *websocket.Conn) *WebsocketConn {
	conn.EnableWriteCompression(true)
	return &WebsocketConn{
		conn: conn,
	}
}

func (c *WebsocketConn) Read(b []byte) (n int, err error) {
	if len(c.rb) == 0 {
		_, c.rb, err = c.conn.ReadMessage()
	}
	n = copy(b, c.rb)
	c.rb = c.rb[n:]
	return
}

func (c *WebsocketConn) Write(b []byte) (n int, err error) {
	err = c.conn.WriteMessage(websocket.BinaryMessage, b)
	n = len(b)
	return
}

func (c *WebsocketConn) Close() error {
	return c.conn.Close()
}

func (c *WebsocketConn) LocalAddr() net.Addr {
	return c.conn.LocalAddr()
}

func (c *WebsocketConn) RemoteAddr() net.Addr {
	return c.conn.RemoteAddr()
}

func (conn *WebsocketConn) SetDeadline(t time.Time) error {
	if err := conn.SetReadDeadline(t); err != nil {
		return err
	}
	return conn.SetWriteDeadline(t)
}
func (c *WebsocketConn) SetReadDeadline(t time.Time) error {
	return c.conn.SetReadDeadline(t)
}

func (c *WebsocketConn) SetWriteDeadline(t time.Time) error {
	return c.conn.SetWriteDeadline(t)
}
<|endoftext|>