cmd/upspin-audit: write collected data to files

Change-Id: I6534220dea5ae08defcc65f76daf4cfde2580dd9
Reviewed-on: https://upspin-review.googlesource.com/17342
Reviewed-by: Rob Pike <r@golang.org>
diff --git a/cmd/upspin-audit/main.go b/cmd/upspin-audit/main.go
index 5d5224d..9e498cc 100644
--- a/cmd/upspin-audit/main.go
+++ b/cmd/upspin-audit/main.go
@@ -8,16 +8,19 @@
 package main
 
 import (
+	"bufio"
 	"flag"
 	"fmt"
 	"log"
 	"os"
 	"path/filepath"
+	"sort"
 
 	"upspin.io/config"
 	"upspin.io/flags"
 	"upspin.io/subcmd"
 	"upspin.io/transports"
+	"upspin.io/upspin"
 	"upspin.io/version"
 )
 
@@ -125,3 +128,27 @@
 	}
 	return fmt.Sprintf("%.2fB", b)
 }
+
+// writeItems sorts and writes a list of reference/size pairs to file.
+func (s *State) writeItems(file string, items []upspin.ListRefsItem) {
+	sort.Slice(items, func(i, j int) bool { return items[i].Ref < items[j].Ref })
+
+	f, err := os.Create(file)
+	if err != nil {
+		s.Exit(err)
+	}
+	defer func() {
+		if err := f.Close(); err != nil {
+			s.Exit(err)
+		}
+	}()
+	w := bufio.NewWriter(f)
+	for _, ri := range items {
+		if _, err := fmt.Fprintf(w, "%q %d\n", ri.Ref, ri.Size); err != nil {
+			s.Exit(err)
+		}
+	}
+	if err := w.Flush(); err != nil {
+		s.Exit(err)
+	}
+}
diff --git a/cmd/upspin-audit/scandir.go b/cmd/upspin-audit/scandir.go
index 538ff11..5feab48 100644
--- a/cmd/upspin-audit/scandir.go
+++ b/cmd/upspin-audit/scandir.go
@@ -9,7 +9,9 @@
 	"flag"
 	"fmt"
 	"os"
+	"path/filepath"
 	"sync"
+	"time"
 
 	"upspin.io/path"
 	"upspin.io/upspin"
@@ -35,6 +37,17 @@
 	done     chan *upspin.DirEntry // Send entries here once it is completely done, including children.
 }
 
+type sizeMap map[upspin.Endpoint]map[upspin.Reference]int64
+
+func (m sizeMap) addRef(ep upspin.Endpoint, ref upspin.Reference, size int64) {
+	refs := m[ep]
+	if refs == nil {
+		refs = make(map[upspin.Reference]int64)
+		m[ep] = refs
+	}
+	refs[ref] = size
+}
+
 func (s *State) scanDirectories(args []string) {
 	const help = `
 Audit scandir scans the directory tree for the named user roots.
@@ -43,7 +56,6 @@
 	fs := flag.NewFlagSet("scandir", flag.ExitOnError)
 	glob := fs.Bool("glob", true, "apply glob processing to the arguments")
 	dataDir := dataDirFlag(fs)
-	_ = dataDir // TODO
 	s.ParseFlags(fs, args, help, "audit scandir root ...")
 
 	if fs.NArg() == 0 || fs.Arg(0) == "help" {
@@ -51,6 +63,10 @@
 		os.Exit(2)
 	}
 
+	if err := os.MkdirAll(*dataDir, 0600); err != nil {
+		s.Exit(err)
+	}
+
 	var paths []upspin.PathName
 	if *glob {
 		paths = s.GlobAllUpspinPath(fs.Args())
@@ -71,6 +87,8 @@
 		}
 	}
 
+	now := time.Now()
+
 	sc := dirScanner{
 		State:    s,
 		buffer:   make(chan *upspin.DirEntry),
@@ -85,11 +103,9 @@
 
 	// Prime the pump.
 	for _, p := range paths {
-		dir := sc.State.DirServer(p)
-		de, err := dir.Lookup(p)
+		de, err := s.DirServer(p).Lookup(p)
 		if err != nil {
-			sc.State.Fail(err)
-			continue
+			s.Exit(err)
 		}
 		sc.do(de)
 	}
@@ -101,17 +117,24 @@
 		close(sc.done)
 	}()
 
-	// Receive the data.
-	size := make(map[upspin.Endpoint]map[upspin.Reference]int64)
+	// Receive and collect the data.
+	size := make(sizeMap)
+	users := make(map[upspin.UserName]sizeMap)
 	for de := range sc.done {
+		p, err := path.Parse(de.Name)
+		if err != nil {
+			s.Fail(err)
+			continue
+		}
+		userSize := users[p.User()]
+		if userSize == nil {
+			userSize = make(sizeMap)
+			users[p.User()] = userSize
+		}
 		for _, block := range de.Blocks {
 			ep := block.Location.Endpoint
-			refs := size[ep]
-			if refs == nil {
-				refs = make(map[upspin.Reference]int64)
-				size[ep] = refs
-			}
-			refs[block.Location.Reference] = block.Size
+			size.addRef(ep, block.Location.Reference, block.Size)
+			userSize.addRef(ep, block.Location.Reference, block.Size)
 		}
 	}
 
@@ -123,11 +146,23 @@
 			sum += s
 		}
 		total += sum
-		fmt.Printf("%s: %d bytes (%s) (%d references)\n", ep, sum, ByteSize(sum), len(refs))
+		fmt.Printf("%s: %d bytes (%s) (%d references)\n", ep.NetAddr, sum, ByteSize(sum), len(refs))
 	}
 	if len(size) > 1 {
 		fmt.Printf("%d bytes total (%s)\n", total, ByteSize(total))
 	}
+
+	// Write the data to files, one for each user/endpoint combo.
+	for u, size := range users {
+		for ep, refs := range size {
+			var items []upspin.ListRefsItem
+			for ref, n := range refs {
+				items = append(items, upspin.ListRefsItem{Ref: ref, Size: n})
+			}
+			file := filepath.Join(*dataDir, fmt.Sprintf("dir.%s.%s.%d", ep.NetAddr, u, now.Unix()))
+			s.writeItems(file, items)
+		}
+	}
 }
 
 // do processes a DirEntry. If it's a file, we deliver it to the done channel.
diff --git a/cmd/upspin-audit/scanstore.go b/cmd/upspin-audit/scanstore.go
index 5a1dd6f..7a6af5d 100644
--- a/cmd/upspin-audit/scanstore.go
+++ b/cmd/upspin-audit/scanstore.go
@@ -9,6 +9,8 @@
 	"flag"
 	"fmt"
 	"os"
+	"path/filepath"
+	"time"
 
 	"upspin.io/bind"
 	"upspin.io/upspin"
@@ -27,7 +29,6 @@
 	fs := flag.NewFlagSet("scanstore", flag.ExitOnError)
 	endpointFlag := fs.String("endpoint", string(s.Config.StoreEndpoint().NetAddr), "network `address` of storage server; default is from config")
 	dataDir := dataDirFlag(fs)
-	_ = dataDir // TODO
 	s.ParseFlags(fs, args, help, "audit scanstore [-endpoint <storeserver address>]")
 
 	if fs.NArg() != 0 { // "audit scanstore help" is covered by this.
@@ -35,19 +36,27 @@
 		os.Exit(2)
 	}
 
+	if err := os.MkdirAll(*dataDir, 0600); err != nil {
+		s.Exit(err)
+	}
+
 	endpoint, err := upspin.ParseEndpoint("remote," + *endpointFlag)
 	if err != nil {
 		s.Exit(err)
 	}
 
+	now := time.Now()
+
 	store, err := bind.StoreServer(s.Config, *endpoint)
 	if err != nil {
 		s.Fail(err)
 		return
 	}
-	token := ""
-	sum := int64(0)
-	numRefs := 0
+	var (
+		token string
+		sum   int64
+		items []upspin.ListRefsItem
+	)
 	for {
 		b, _, _, err := store.Get(upspin.ListRefsMetadata + upspin.Reference(token))
 		if err != nil {
@@ -60,14 +69,16 @@
 			s.Exit(err)
 			return
 		}
-		for _, r := range refs.Refs {
-			numRefs++
-			sum += r.Size
+		for _, ri := range refs.Refs {
+			sum += ri.Size
+			items = append(items, ri)
 		}
 		token = refs.Next
 		if token == "" {
 			break
 		}
 	}
-	fmt.Printf("%s: %d bytes total (%s) in %d references\n", endpoint, sum, ByteSize(sum), numRefs)
+	fmt.Printf("%s: %d bytes total (%s) in %d references\n", endpoint.NetAddr, sum, ByteSize(sum), len(items))
+	file := filepath.Join(*dataDir, fmt.Sprintf("store.%s.%d", endpoint.NetAddr, now.Unix()))
+	s.writeItems(file, items)
 }