Golang KeepClient.SetServiceRoots方法代码示例

// PullItemAndProcess pulls items from PullQueue and processes them.
//	For each Pull request:
//		Generate a random API token.
//		Generate a permission signature using this token, timestamp ~60 seconds in the future, and desired block hash.
//		Using this token & signature, retrieve the given block.
//		Write to storage
//
func PullItemAndProcess(pullRequest PullRequest, token string, keepClient *keepclient.KeepClient) (err error) {
	keepClient.Arvados.ApiToken = token

	serviceRoots := make(map[string]string)
	for _, addr := range pullRequest.Servers {
		serviceRoots[addr] = addr
	}
	keepClient.SetServiceRoots(serviceRoots, nil, nil)

	// Generate signature with a random token
	expiresAt := time.Now().Add(60 * time.Second)
	signedLocator := SignLocator(pullRequest.Locator, token, expiresAt)

	reader, contentLen, _, err := GetContent(signedLocator, keepClient)
	if err != nil {
		return
	}
	if reader == nil {
		return errors.New(fmt.Sprintf("No reader found for : %s", signedLocator))
	}
	defer reader.Close()

	readContent, err := ioutil.ReadAll(reader)
	if err != nil {
		return err
	}

	if (readContent == nil) || (int64(len(readContent)) != contentLen) {
		return errors.New(fmt.Sprintf("Content not found for: %s", signedLocator))
	}

	err = PutContent(readContent, pullRequest.Locator)
	return
}

func sendTrashListError(c *C, server *httptest.Server) {
	tl := map[string]TrashList{
		server.URL: TrashList{TrashRequest{"000000000000000000000000deadbeef", 99}}}

	kc := keepclient.KeepClient{Client: &http.Client{}}
	kc.SetServiceRoots(map[string]string{"xxxx": server.URL},
		map[string]string{"xxxx": server.URL},
		map[string]string{})

	err := SendTrashLists("", &kc, tl)

	c.Check(err, NotNil)
	c.Check(err[0], NotNil)
}

func sendTrashListError(c *C, server *httptest.Server) {
	tl := map[string]TrashList{
		server.URL: TrashList{TrashRequest{"000000000000000000000000deadbeef", 99}}}

	arv := arvadosclient.ArvadosClient{ApiToken: "abc123"}
	kc := keepclient.KeepClient{Arvados: &arv, Client: &http.Client{}}
	kc.SetServiceRoots(map[string]string{"xxxx": server.URL},
		map[string]string{"xxxx": server.URL},
		map[string]string{})

	err := SendTrashLists(nil, &kc, tl, false)

	c.Check(err, NotNil)
	c.Check(err[0], NotNil)
}

func overrideServices(kc *keepclient.KeepClient) {
	roots := make(map[string]string)
	if *ServiceURL != "" {
		roots["zzzzz-bi6l4-000000000000000"] = *ServiceURL
	} else if *ServiceUUID != "" {
		for uuid, url := range kc.GatewayRoots() {
			if uuid == *ServiceUUID {
				roots[uuid] = url
				break
			}
		}
		if len(roots) == 0 {
			log.Fatalf("Service %q was not in list advertised by API %+q", *ServiceUUID, kc.GatewayRoots())
		}
	} else {
		return
	}
	kc.SetServiceRoots(roots, roots, roots)
}

func testGetKeepServersFromAPI(c *C, testData APITestData, expectedError string) {
	keepServers := ServiceList{
		ItemsAvailable: testData.numServers,
		KeepServers: []ServerAddress{{
			SSL:         false,
			Host:        "example.com",
			Port:        12345,
			UUID:        "abcdefg",
			ServiceType: testData.serverType,
		}},
	}

	ksJSON, _ := json.Marshal(keepServers)
	apiStubResponses := make(map[string]arvadostest.StubResponse)
	apiStubResponses["/arvados/v1/keep_services"] = arvadostest.StubResponse{testData.statusCode, string(ksJSON)}
	apiStub := arvadostest.ServerStub{apiStubResponses}

	api := httptest.NewServer(&apiStub)
	defer api.Close()

	arv := arvadosclient.ArvadosClient{
		Scheme:    "http",
		ApiServer: api.URL[7:],
		ApiToken:  "abc123",
		Client:    &http.Client{Transport: &http.Transport{}},
	}

	kc := keepclient.KeepClient{Arvados: &arv, Client: &http.Client{}}
	kc.SetServiceRoots(map[string]string{"xxxx": "http://example.com:23456"},
		map[string]string{"xxxx": "http://example.com:23456"},
		map[string]string{})

	params := GetKeepServersParams{
		Client: arv,
		Logger: nil,
		Limit:  10,
	}

	_, err := GetKeepServersAndSummarize(params)
	c.Assert(err, NotNil)
	c.Assert(err, ErrorMatches, fmt.Sprintf(".*%s.*", expectedError))
}

func (s *KeepSuite) TestSendTrashLists(c *C) {
	th := TestHandler{}
	server := httptest.NewServer(&th)

	tl := map[string]TrashList{
		server.URL: TrashList{TrashRequest{"000000000000000000000000deadbeef", 99}}}

	kc := keepclient.KeepClient{Client: &http.Client{}}
	kc.SetServiceRoots(map[string]string{"xxxx": server.URL},
		map[string]string{"xxxx": server.URL},
		map[string]string{})

	err := SendTrashLists("", &kc, tl)
	server.Close()

	c.Check(err, IsNil)

	c.Check(th.request,
		DeepEquals,
		tl[server.URL])

}

func (s *KeepSuite) TestSendTrashLists(c *C) {
	th := TestHandler{}
	server := httptest.NewServer(&th)
	defer server.Close()

	tl := map[string]TrashList{
		server.URL: TrashList{TrashRequest{"000000000000000000000000deadbeef", 99}}}

	arv := arvadosclient.ArvadosClient{ApiToken: "abc123"}
	kc := keepclient.KeepClient{Arvados: &arv, Client: &http.Client{}}
	kc.SetServiceRoots(map[string]string{"xxxx": server.URL},
		map[string]string{"xxxx": server.URL},
		map[string]string{})

	err := SendTrashLists(nil, &kc, tl, false)

	c.Check(err, IsNil)

	c.Check(th.request,
		DeepEquals,
		tl[server.URL])

}

func runProxy(c *C, args []string, port int, bogusClientToken bool) keepclient.KeepClient {
	if bogusClientToken {
		os.Setenv("ARVADOS_API_TOKEN", "bogus-token")
	}
	arv, err := arvadosclient.MakeArvadosClient()
	c.Assert(err, Equals, nil)
	kc := keepclient.KeepClient{
		Arvados:       &arv,
		Want_replicas: 2,
		Using_proxy:   true,
		Client:        &http.Client{},
	}
	locals := map[string]string{
		"proxy": fmt.Sprintf("http://localhost:%v", port),
	}
	writableLocals := map[string]string{
		"proxy": fmt.Sprintf("http://localhost:%v", port),
	}
	kc.SetServiceRoots(locals, writableLocals, nil)
	c.Check(kc.Using_proxy, Equals, true)
	c.Check(len(kc.LocalRoots()), Equals, 1)
	for _, root := range kc.LocalRoots() {
		c.Check(root, Equals, fmt.Sprintf("http://localhost:%v", port))
	}
	log.Print("keepclient created")
	if bogusClientToken {
		arvadostest.ResetEnv()
	}

	{
		os.Args = append(args, fmt.Sprintf("-listen=:%v", port))
		listener = nil
		go main()
	}

	return kc
}

func testGetKeepServersAndSummarize(c *C, testData KeepServerTestData) {
	ksStubResponses := make(map[string]arvadostest.StubResponse)
	ksStubResponses["/status.json"] = arvadostest.StubResponse{testData.statusStatusCode, string(`{}`)}
	ksStubResponses["/index"] = arvadostest.StubResponse{testData.indexStatusCode, testData.indexResponseBody}
	ksStub := arvadostest.ServerStub{ksStubResponses}
	ks := httptest.NewServer(&ksStub)
	defer ks.Close()

	ksURL, err := url.Parse(ks.URL)
	c.Check(err, IsNil)
	ksHost, port, err := net.SplitHostPort(ksURL.Host)
	ksPort, err := strconv.Atoi(port)
	c.Check(err, IsNil)

	servers_list := ServiceList{
		ItemsAvailable: 1,
		KeepServers: []ServerAddress{{
			SSL:         false,
			Host:        ksHost,
			Port:        ksPort,
			UUID:        "abcdefg",
			ServiceType: "disk",
		}},
	}
	ksJSON, _ := json.Marshal(servers_list)
	apiStubResponses := make(map[string]arvadostest.StubResponse)
	apiStubResponses["/arvados/v1/keep_services"] = arvadostest.StubResponse{200, string(ksJSON)}
	apiStub := arvadostest.ServerStub{apiStubResponses}

	api := httptest.NewServer(&apiStub)
	defer api.Close()

	arv := arvadosclient.ArvadosClient{
		Scheme:    "http",
		ApiServer: api.URL[7:],
		ApiToken:  "abc123",
		Client:    &http.Client{Transport: &http.Transport{}},
	}

	kc := keepclient.KeepClient{Arvados: &arv, Client: &http.Client{}}
	kc.SetServiceRoots(map[string]string{"xxxx": ks.URL},
		map[string]string{"xxxx": ks.URL},
		map[string]string{})

	params := GetKeepServersParams{
		Client: arv,
		Logger: nil,
		Limit:  10,
	}

	// GetKeepServersAndSummarize
	results, err := GetKeepServersAndSummarize(params)

	if testData.expectedError == "" {
		c.Assert(err, IsNil)
		c.Assert(results, NotNil)

		blockToServers := results.BlockToServers

		blockLocators := strings.Split(testData.indexResponseBody, "\n")
		for _, loc := range blockLocators {
			locator := strings.Split(loc, " ")[0]
			if locator != "" {
				blockLocator, err := blockdigest.ParseBlockLocator(locator)
				c.Assert(err, IsNil)

				blockDigestWithSize := blockdigest.DigestWithSize{blockLocator.Digest, uint32(blockLocator.Size)}
				blockServerInfo := blockToServers[blockDigestWithSize]
				c.Assert(blockServerInfo[0].Mtime, NotNil)
			}
		}
	} else {
		c.Assert(err, ErrorMatches, testData.expectedError)
	}
}