func (s *ServerRequiredSuite) TestPutAskGet(c *C) {
log.Print("TestPutAndGet start")
os.Args = []string{"keepproxy", "-listen=:29950"}
os.Setenv("ARVADOS_API_TOKEN", "4axaw8zxe0qm22wa6urpp5nskcne8z88cvbupv653y1njyi05h")
go main()
time.Sleep(100 * time.Millisecond)
setupProxyService()
os.Setenv("ARVADOS_EXTERNAL_CLIENT", "true")
arv, err := sdk.MakeArvadosClient()
kc, err := keepclient.MakeKeepClient(&arv)
c.Check(kc.Arvados.External, Equals, true)
c.Check(kc.Using_proxy, Equals, true)
c.Check(len(kc.ServiceRoots()), Equals, 1)
c.Check(kc.ServiceRoots()[0], Equals, "http://localhost:29950")
c.Check(err, Equals, nil)
os.Setenv("ARVADOS_EXTERNAL_CLIENT", "")
log.Print("keepclient created")
defer listener.Close()
hash := fmt.Sprintf("%x", md5.Sum([]byte("foo")))
var hash2 string
{
_, _, err := kc.Ask(hash)
c.Check(err, Equals, keepclient.BlockNotFound)
log.Print("Ask 1")
}
{
var rep int
var err error
hash2, rep, err = kc.PutB([]byte("foo"))
c.Check(hash2, Equals, fmt.Sprintf("%s+3", hash))
c.Check(rep, Equals, 2)
c.Check(err, Equals, nil)
log.Print("PutB")
}
{
blocklen, _, err := kc.Ask(hash2)
c.Assert(err, Equals, nil)
c.Check(blocklen, Equals, int64(3))
log.Print("Ask 2")
}
{
reader, blocklen, _, err := kc.Get(hash2)
c.Assert(err, Equals, nil)
all, err := ioutil.ReadAll(reader)
c.Check(all, DeepEquals, []byte("foo"))
c.Check(blocklen, Equals, int64(3))
log.Print("Get")
}
log.Print("TestPutAndGet done")
}
func (s *StandaloneSuite) TestPutProxyInsufficientReplicas(c *C) {
log.Printf("TestPutProxy")
st := StubProxyHandler{make(chan string, 1)}
arv, err := sdk.MakeArvadosClient()
kc, _ := MakeKeepClient(&arv)
kc.Want_replicas = 3
kc.Using_proxy = true
arv.ApiToken = "abc123"
service_roots := make([]string, 1)
ks1 := RunSomeFakeKeepServers(st, 1, 2990)
for i, k := range ks1 {
service_roots[i] = k.url
defer k.listener.Close()
}
kc.SetServiceRoots(service_roots)
_, replicas, err := kc.PutB([]byte("foo"))
<-st.handled
c.Check(err, Equals, InsufficientReplicasError)
c.Check(replicas, Equals, 2)
log.Printf("TestPutProxy done")
}
func (s *StandaloneSuite) TestChecksum(c *C) {
foohash := fmt.Sprintf("%x", md5.Sum([]byte("foo")))
barhash := fmt.Sprintf("%x", md5.Sum([]byte("bar")))
st := BarHandler{make(chan string, 1)}
listener, url := RunBogusKeepServer(st, 2990)
defer listener.Close()
arv, err := sdk.MakeArvadosClient()
kc, _ := MakeKeepClient(&arv)
arv.ApiToken = "abc123"
kc.SetServiceRoots([]string{url})
r, n, _, err := kc.Get(barhash)
_, err = ioutil.ReadAll(r)
c.Check(n, Equals, int64(3))
c.Check(err, Equals, nil)
<-st.handled
r, n, _, err = kc.Get(foohash)
_, err = ioutil.ReadAll(r)
c.Check(n, Equals, int64(3))
c.Check(err, Equals, BadChecksum)
<-st.handled
}
func (s *StandaloneSuite) TestGet(c *C) {
log.Printf("TestGet")
hash := fmt.Sprintf("%x", md5.Sum([]byte("foo")))
st := StubGetHandler{
c,
hash,
"abc123",
[]byte("foo")}
listener, url := RunBogusKeepServer(st, 2990)
defer listener.Close()
arv, err := sdk.MakeArvadosClient()
kc, _ := MakeKeepClient(&arv)
arv.ApiToken = "abc123"
kc.SetServiceRoots([]string{url})
r, n, url2, err := kc.Get(hash)
defer r.Close()
c.Check(err, Equals, nil)
c.Check(n, Equals, int64(3))
c.Check(url2, Equals, fmt.Sprintf("%s/%s", url, hash))
content, err2 := ioutil.ReadAll(r)
c.Check(err2, Equals, nil)
c.Check(content, DeepEquals, []byte("foo"))
log.Printf("TestGet done")
}
func (s *StandaloneSuite) TestPutHR(c *C) {
log.Printf("TestPutHR")
hash := fmt.Sprintf("%x", md5.Sum([]byte("foo")))
st := StubPutHandler{
c,
hash,
"abc123",
"foo",
make(chan string, 2)}
arv, _ := sdk.MakeArvadosClient()
kc, _ := MakeKeepClient(&arv)
kc.Want_replicas = 2
arv.ApiToken = "abc123"
service_roots := make([]string, 5)
ks := RunSomeFakeKeepServers(st, 5, 2990)
for i := 0; i < len(ks); i += 1 {
service_roots[i] = ks[i].url
defer ks[i].listener.Close()
}
kc.SetServiceRoots(service_roots)
reader, writer := io.Pipe()
go func() {
writer.Write([]byte("foo"))
writer.Close()
}()
kc.PutHR(hash, reader, 3)
shuff := kc.shuffledServiceRoots(hash)
log.Print(shuff)
s1 := <-st.handled
s2 := <-st.handled
c.Check((s1 == shuff[0] && s2 == shuff[1]) ||
(s1 == shuff[1] && s2 == shuff[0]),
Equals,
true)
log.Printf("TestPutHR done")
}
func (s *ServerRequiredSuite) TestMakeKeepClient(c *C) {
os.Setenv("ARVADOS_API_HOST", "localhost:3001")
os.Setenv("ARVADOS_API_TOKEN", "4axaw8zxe0qm22wa6urpp5nskcne8z88cvbupv653y1njyi05h")
os.Setenv("ARVADOS_API_HOST_INSECURE", "true")
arv, err := sdk.MakeArvadosClient()
c.Assert(err, Equals, nil)
kc, err := MakeKeepClient(&arv)
c.Assert(err, Equals, nil)
c.Check(len(kc.ServiceRoots()), Equals, 2)
c.Check(kc.ServiceRoots()[0], Equals, "http://localhost:25107")
c.Check(kc.ServiceRoots()[1], Equals, "http://localhost:25108")
}
func (s *StandaloneSuite) TestPutWithFail(c *C) {
log.Printf("TestPutWithFail")
hash := fmt.Sprintf("%x", md5.Sum([]byte("foo")))
st := StubPutHandler{
c,
hash,
"abc123",
"foo",
make(chan string, 2)}
fh := FailHandler{
make(chan string, 1)}
arv, err := sdk.MakeArvadosClient()
kc, _ := MakeKeepClient(&arv)
kc.Want_replicas = 2
arv.ApiToken = "abc123"
service_roots := make([]string, 5)
ks1 := RunSomeFakeKeepServers(st, 4, 2990)
ks2 := RunSomeFakeKeepServers(fh, 1, 2995)
for i, k := range ks1 {
service_roots[i] = k.url
defer k.listener.Close()
}
for i, k := range ks2 {
service_roots[len(ks1)+i] = k.url
defer k.listener.Close()
}
kc.SetServiceRoots(service_roots)
shuff := kc.shuffledServiceRoots(fmt.Sprintf("%x", md5.Sum([]byte("foo"))))
phash, replicas, err := kc.PutB([]byte("foo"))
<-fh.handled
c.Check(err, Equals, nil)
c.Check(phash, Equals, "")
c.Check(replicas, Equals, 2)
c.Check(<-st.handled, Equals, shuff[1])
c.Check(<-st.handled, Equals, shuff[2])
}
func UploadToStubHelper(c *C, st http.Handler, f func(KeepClient, string,
io.ReadCloser, io.WriteCloser, chan uploadStatus)) {
listener, url := RunBogusKeepServer(st, 2990)
defer listener.Close()
arv, _ := sdk.MakeArvadosClient()
arv.ApiToken = "abc123"
kc, _ := MakeKeepClient(&arv)
reader, writer := io.Pipe()
upload_status := make(chan uploadStatus)
f(kc, url, reader, writer, upload_status)
}
func runProxy(c *C, args []string, token string, port int) keepclient.KeepClient {
os.Args = append(args, fmt.Sprintf("-listen=:%v", port))
os.Setenv("ARVADOS_API_TOKEN", "4axaw8zxe0qm22wa6urpp5nskcne8z88cvbupv653y1njyi05h")
go main()
time.Sleep(100 * time.Millisecond)
os.Setenv("ARVADOS_KEEP_PROXY", fmt.Sprintf("http://localhost:%v", port))
os.Setenv("ARVADOS_API_TOKEN", token)
arv, err := sdk.MakeArvadosClient()
kc, err := keepclient.MakeKeepClient(&arv)
c.Check(kc.Using_proxy, Equals, true)
c.Check(len(kc.ServiceRoots()), Equals, 1)
c.Check(kc.ServiceRoots()[0], Equals, fmt.Sprintf("http://localhost:%v", port))
c.Check(err, Equals, nil)
os.Setenv("ARVADOS_KEEP_PROXY", "")
log.Print("keepclient created")
return kc
}
func (s *StandaloneSuite) TestGetFail(c *C) {
hash := fmt.Sprintf("%x", md5.Sum([]byte("foo")))
st := FailHandler{make(chan string, 1)}
listener, url := RunBogusKeepServer(st, 2990)
defer listener.Close()
arv, err := sdk.MakeArvadosClient()
kc, _ := MakeKeepClient(&arv)
arv.ApiToken = "abc123"
kc.SetServiceRoots([]string{url})
r, n, url2, err := kc.Get(hash)
c.Check(err, Equals, BlockNotFound)
c.Check(n, Equals, int64(0))
c.Check(url2, Equals, "")
c.Check(r, Equals, nil)
}
func (s *StandaloneSuite) TestGetWithFailures(c *C) {
hash := fmt.Sprintf("%x", md5.Sum([]byte("foo")))
fh := FailHandler{
make(chan string, 1)}
st := StubGetHandler{
c,
hash,
"abc123",
[]byte("foo")}
arv, err := sdk.MakeArvadosClient()
kc, _ := MakeKeepClient(&arv)
arv.ApiToken = "abc123"
service_roots := make([]string, 5)
ks1 := RunSomeFakeKeepServers(st, 1, 2990)
ks2 := RunSomeFakeKeepServers(fh, 4, 2991)
for i, k := range ks1 {
service_roots[i] = k.url
defer k.listener.Close()
}
for i, k := range ks2 {
service_roots[len(ks1)+i] = k.url
defer k.listener.Close()
}
kc.SetServiceRoots(service_roots)
r, n, url2, err := kc.Get(hash)
<-fh.handled
c.Check(err, Equals, nil)
c.Check(n, Equals, int64(3))
c.Check(url2, Equals, fmt.Sprintf("%s/%s", ks1[0].url, hash))
content, err2 := ioutil.ReadAll(r)
c.Check(err2, Equals, nil)
c.Check(content, DeepEquals, []byte("foo"))
}
func (s *ServerRequiredSuite) TestPutGetHead(c *C) {
os.Setenv("ARVADOS_API_HOST", "localhost:3001")
os.Setenv("ARVADOS_API_TOKEN", "4axaw8zxe0qm22wa6urpp5nskcne8z88cvbupv653y1njyi05h")
os.Setenv("ARVADOS_API_HOST_INSECURE", "true")
arv, err := sdk.MakeArvadosClient()
kc, err := MakeKeepClient(&arv)
c.Assert(err, Equals, nil)
hash := fmt.Sprintf("%x", md5.Sum([]byte("foo")))
{
n, _, err := kc.Ask(hash)
c.Check(err, Equals, BlockNotFound)
c.Check(n, Equals, int64(0))
}
{
hash2, replicas, err := kc.PutB([]byte("foo"))
c.Check(hash2, Equals, fmt.Sprintf("%s+%v", hash, 3))
c.Check(replicas, Equals, 2)
c.Check(err, Equals, nil)
}
{
r, n, url2, err := kc.Get(hash)
c.Check(err, Equals, nil)
c.Check(n, Equals, int64(3))
c.Check(url2, Equals, fmt.Sprintf("http://localhost:25108/%s", hash))
content, err2 := ioutil.ReadAll(r)
c.Check(err2, Equals, nil)
c.Check(content, DeepEquals, []byte("foo"))
}
{
n, url2, err := kc.Ask(hash)
c.Check(err, Equals, nil)
c.Check(n, Equals, int64(3))
c.Check(url2, Equals, fmt.Sprintf("http://localhost:25108/%s", hash))
}
}
func main() {
var (
listen string
no_get bool
no_put bool
default_replicas int
pidfile string
)
flagset := flag.NewFlagSet("default", flag.ExitOnError)
flagset.StringVar(
&listen,
"listen",
DEFAULT_ADDR,
"Interface on which to listen for requests, in the format "+
"ipaddr:port. e.g. -listen=10.0.1.24:8000. Use -listen=:port "+
"to listen on all network interfaces.")
flagset.BoolVar(
&no_get,
"no-get",
false,
"If set, disable GET operations")
flagset.BoolVar(
&no_put,
"no-put",
false,
"If set, disable PUT operations")
flagset.IntVar(
&default_replicas,
"default-replicas",
2,
"Default number of replicas to write if not specified by the client.")
flagset.StringVar(
&pidfile,
"pid",
"",
"Path to write pid file")
flagset.Parse(os.Args[1:])
arv, err := sdk.MakeArvadosClient()
if err != nil {
log.Fatalf("Error setting up arvados client %s", err.Error())
}
kc, err := keepclient.MakeKeepClient(&arv)
if err != nil {
log.Fatalf("Error setting up keep client %s", err.Error())
}
if pidfile != "" {
f, err := os.Create(pidfile)
if err == nil {
fmt.Fprint(f, os.Getpid())
f.Close()
} else {
log.Printf("Error writing pid file (%s): %s", pidfile, err.Error())
}
}
kc.Want_replicas = default_replicas
listener, err = net.Listen("tcp", listen)
if err != nil {
log.Fatalf("Could not listen on %v", listen)
}
go RefreshServicesList(&kc)
// Shut down the server gracefully (by closing the listener)
// if SIGTERM is received.
term := make(chan os.Signal, 1)
go func(sig <-chan os.Signal) {
s := <-sig
log.Println("caught signal:", s)
listener.Close()
}(term)
signal.Notify(term, syscall.SIGTERM)
signal.Notify(term, syscall.SIGINT)
if pidfile != "" {
f, err := os.Create(pidfile)
if err == nil {
fmt.Fprint(f, os.Getpid())
f.Close()
} else {
log.Printf("Error writing pid file (%s): %s", pidfile, err.Error())
}
}
log.Printf("Arvados Keep proxy started listening on %v with server list %v", listener.Addr(), kc.ServiceRoots())
// Start listening for requests.
http.Serve(listener, MakeRESTRouter(!no_get, !no_put, &kc))
log.Println("shutting down")
if pidfile != "" {
os.Remove(pidfile)
}
}