Golang keepclient.KeepClient類代碼示例

// Get block headers from keep. Log any errors.
func performKeepBlockCheck(kc *keepclient.KeepClient, blobSignatureTTL time.Duration, blobSigningKey string, blockLocators []string, verbose bool) error {
	totalBlocks := len(blockLocators)
	notFoundBlocks := 0
	current := 0
	for _, locator := range blockLocators {
		current++
		if verbose {
			log.Printf("Verifying block %d of %d: %v", current, totalBlocks, locator)
		}
		getLocator := locator
		if blobSigningKey != "" {
			expiresAt := time.Now().AddDate(0, 0, 1)
			getLocator = keepclient.SignLocator(locator, kc.Arvados.ApiToken, expiresAt, blobSignatureTTL, []byte(blobSigningKey))
		}

		_, _, err := kc.Ask(getLocator)
		if err != nil {
			notFoundBlocks++
			log.Printf("Error verifying block %v: %v", locator, err)
		}
	}

	log.Printf("Verify block totals: %d attempts, %d successes, %d errors", totalBlocks, totalBlocks-notFoundBlocks, notFoundBlocks)

	if notFoundBlocks > 0 {
		return fmt.Errorf("Block verification failed for %d out of %d blocks with matching prefix.", notFoundBlocks, totalBlocks)
	}

	return nil
}

// Copy blocks from src to dst; only those that are missing in dst are copied
func copyBlocksToDst(toBeCopied []string, kcSrc, kcDst *keepclient.KeepClient, blobSigningKey string) error {
	total := len(toBeCopied)

	startedAt := time.Now()
	for done, locator := range toBeCopied {
		if done == 0 {
			log.Printf("Copying data block %d of %d (%.2f%% done): %v", done+1, total,
				float64(done)/float64(total)*100, locator)
		} else {
			timePerBlock := time.Since(startedAt) / time.Duration(done)
			log.Printf("Copying data block %d of %d (%.2f%% done, %v est. time remaining): %v", done+1, total,
				float64(done)/float64(total)*100, timePerBlock*time.Duration(total-done), locator)
		}

		getLocator := locator
		expiresAt := time.Now().AddDate(0, 0, 1)
		if blobSigningKey != "" {
			getLocator = keepclient.SignLocator(getLocator, kcSrc.Arvados.ApiToken, expiresAt, []byte(blobSigningKey))
		}

		reader, len, _, err := kcSrc.Get(getLocator)
		if err != nil {
			return fmt.Errorf("Error getting block: %v %v", locator, err)
		}

		_, _, err = kcDst.PutHR(getLocator[:32], reader, len)
		if err != nil {
			return fmt.Errorf("Error copying data block: %v %v", locator, err)
		}
	}

	log.Printf("Successfully copied to destination %d blocks.", total)
	return nil
}

// 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
}

// Do a Put in the first and Get from the second,
// which should raise block not found error.
func testNoCrosstalk(c *C, testData string, kc1, kc2 *keepclient.KeepClient) {
	// Put a block using kc1
	locator, _, err := kc1.PutB([]byte(testData))
	c.Assert(err, Equals, nil)

	locator = strings.Split(locator, "+")[0]
	_, _, _, err = kc2.Get(keepclient.SignLocator(locator, kc2.Arvados.ApiToken, time.Now().AddDate(0, 0, 1), blobSignatureTTL, []byte(blobSigningKey)))
	c.Assert(err, NotNil)
	c.Check(err.Error(), Equals, "Block not found")
}

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 doWrites(kc *keepclient.KeepClient, nextBuf chan []byte, nextLocator chan string) {
	for buf := range nextBuf {
		locator, _, err := kc.PutB(buf)
		if err != nil {
			log.Print(err)
			errorsChan <- struct{}{}
			continue
		}
		bytesOutChan <- uint64(len(buf))
		for cap(nextLocator) > len(nextLocator)+*WriteThreads {
			// Give the readers something to do, unless
			// they have lots queued up already.
			nextLocator <- locator
		}
	}
}

// Get list of unique locators from the specified cluster
func getUniqueLocators(kc *keepclient.KeepClient, prefix string) (map[string]bool, error) {
	uniqueLocators := map[string]bool{}

	// Get index and dedup
	for uuid := range kc.LocalRoots() {
		reader, err := kc.GetIndex(uuid, prefix)
		if err != nil {
			return uniqueLocators, err
		}
		scanner := bufio.NewScanner(reader)
		for scanner.Scan() {
			uniqueLocators[strings.Split(scanner.Text(), " ")[0]] = true
		}
	}

	return uniqueLocators, nil
}

// ComputePullServers creates a map from block locator to PullServers
// with one entry for each under-replicated block.
//
// This method ignores zero-replica blocks since there are no servers
// to pull them from, so callers should feel free to omit them, but
// this function will ignore them if they are provided.
func ComputePullServers(kc *keepclient.KeepClient,
	keepServerInfo *keep.ReadServers,
	blockToDesiredReplication map[blockdigest.DigestWithSize]int,
	underReplicated BlockSet) (m map[Locator]PullServers) {
	m = map[Locator]PullServers{}
	// We use CanonicalString to avoid filling memory with dupicate
	// copies of the same string.
	var cs CanonicalString

	// Servers that are writeable
	writableServers := map[string]struct{}{}
	for _, url := range kc.WritableLocalRoots() {
		writableServers[cs.Get(url)] = struct{}{}
	}

	for block := range underReplicated {
		serversStoringBlock := keepServerInfo.BlockToServers[block]
		numCopies := len(serversStoringBlock)
		numCopiesMissing := blockToDesiredReplication[block] - numCopies
		if numCopiesMissing > 0 {
			// We expect this to always be true, since the block was listed
			// in underReplicated.

			if numCopies > 0 {
				// Not much we can do with blocks with no copies.

				// A server's host-port string appears as a key in this map
				// iff it contains the block.
				serverHasBlock := map[string]struct{}{}
				for _, info := range serversStoringBlock {
					sa := keepServerInfo.KeepServerIndexToAddress[info.ServerIndex]
					serverHasBlock[cs.Get(sa.URL())] = struct{}{}
				}

				roots := keepclient.NewRootSorter(kc.LocalRoots(),
					block.String()).GetSortedRoots()

				l := Locator(block)
				m[l] = CreatePullServers(cs, serverHasBlock, writableServers,
					roots, numCopiesMissing)
			}
		}
	}
	return m
}

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 doReads(kc *keepclient.KeepClient, nextLocator chan string) {
	for locator := range nextLocator {
		rdr, size, url, err := kc.Get(locator)
		if err != nil {
			log.Print(err)
			errorsChan <- struct{}{}
			continue
		}
		n, err := io.Copy(ioutil.Discard, rdr)
		rdr.Close()
		if n != size || err != nil {
			log.Printf("Got %d bytes (expected %d) from %s: %v", n, size, url, err)
			errorsChan <- struct{}{}
			continue
			// Note we don't count the bytes received in
			// partial/corrupt responses: we are measuring
			// throughput, not resource consumption.
		}
		bytesInChan <- uint64(n)
	}
}

// BuildTrashLists builds list of blocks to be sent to trash queue
func BuildTrashLists(kc *keepclient.KeepClient,
	keepServerInfo *keep.ReadServers,
	keepBlocksNotInCollections BlockSet) (m map[string]keep.TrashList, err error) {

	// Servers that are writeable
	writableServers := map[string]struct{}{}
	for _, url := range kc.WritableLocalRoots() {
		writableServers[url] = struct{}{}
	}

	_ttl, err := kc.Arvados.Discovery("blobSignatureTtl")
	if err != nil {
		return nil, errors.New(fmt.Sprintf("Failed to get blobSignatureTtl, can't build trash lists: %v", err))
	}

	ttl := int64(_ttl.(float64))

	// expire unreferenced blocks more than "ttl" seconds old.
	expiry := time.Now().UTC().Unix() - ttl

	return buildTrashListsInternal(writableServers, keepServerInfo, expiry, keepBlocksNotInCollections), nil
}

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 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)
}