1// Copyright 2018 The Bazel Authors. All rights reserved. 2// 3// Licensed under the Apache License, Version 2.0 (the "License"); 4// you may not use this file except in compliance with the License. 5// You may obtain a copy of the License at 6// 7// http://www.apache.org/licenses/LICENSE-2.0 8// 9// Unless required by applicable law or agreed to in writing, software 10// distributed under the License is distributed on an "AS IS" BASIS, 11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12// See the License for the specific language governing permissions and 13// limitations under the License. 14 15// Package shard provides functions to help sharding your data. 16package shard 17 18import ( 19 "archive/zip" 20 "errors" 21 "fmt" 22 "hash/fnv" 23 "io" 24 "strings" 25) 26 27// Func converts a name and a number of shards into a particular shard index. 28type Func func(name string, shardCount int) int 29 30// FNV uses the FNV hash algo on the provided string and mods its result by shardCount. 31func FNV(name string, shardCount int) int { 32 h := fnv.New32() 33 h.Write([]byte(name)) 34 return int(h.Sum32()) % shardCount 35} 36 37// MakeSepFunc creates a shard function that takes a substring from 0 to the last occurrence of 38// separator from the name to be sharded, and passes that onto the provided shard function. 39func MakeSepFunc(sep string, s Func) Func { 40 return func(name string, shardCount int) int { 41 idx := strings.LastIndex(name, sep) 42 if idx == -1 { 43 return s(name, shardCount) 44 } 45 return s(name[:idx], shardCount) 46 } 47} 48 49// ZipShard takes a given zip reader, and shards its content across the provided io.Writers 50// utilizing the provided SharderFunc. 51func ZipShard(r *zip.Reader, zws []*zip.Writer, fn Func) error { 52 sc := len(zws) 53 if sc == 0 { 54 return errors.New("no output writers") 55 } 56 57 for _, f := range r.File { 58 if !f.Mode().IsRegular() { 59 continue 60 } 61 si := fn(f.Name, sc) 62 if si < 0 || si > sc { 63 return fmt.Errorf("s.Shard(%s, %d) yields invalid shard index: %d", f.Name, sc, si) 64 } 65 zw := zws[si] 66 var rc io.ReadCloser 67 rc, err := f.Open() 68 if err != nil { 69 return fmt.Errorf("%s: could not open: %v", f.Name, err) 70 } 71 var zo io.Writer 72 zo, err = zw.CreateHeader(&zip.FileHeader{ 73 Name: f.Name, 74 Method: zip.Store, 75 }) 76 if err != nil { 77 return fmt.Errorf("%s: could not create output entry: %v", f.Name, err) 78 } 79 if err := copyAndClose(zo, rc); err != nil { 80 return fmt.Errorf("%s: copy to output failed: %v", f.Name, err) 81 } 82 } 83 return nil 84} 85 86func copyAndClose(w io.Writer, rc io.ReadCloser) error { 87 defer rc.Close() 88 _, err := io.Copy(w, rc) 89 return err 90} 91