1// Copyright 2022 The BoringSSL Authors
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//     https://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
15package runner
16
17import (
18	"fmt"
19	"os"
20	"strconv"
21)
22
23const (
24	shardStatusFileEnv = "TEST_SHARD_STATUS_FILE"
25	shardTotalEnv      = "TEST_TOTAL_SHARDS"
26	shardIndexEnv      = "TEST_SHARD_INDEX"
27	shardPrefix        = "RUNNER_"
28)
29
30func init() {
31	// When run under `go test`, init() functions may be run twice if the
32	// test binary ends up forking and execing itself. Therefore we move
33	// the environment variables to names that don't interfere with Go's
34	// own support for sharding. If we recorded and erased them, then they
35	// wouldn't exist the second time the binary runs.
36	for _, key := range []string{shardStatusFileEnv, shardTotalEnv, shardIndexEnv} {
37		value := os.Getenv(key)
38		if len(value) > 0 {
39			os.Setenv(shardPrefix+key, value)
40			os.Setenv(key, "")
41		}
42	}
43}
44
45// getSharding returns the shard index and count, or zeros if sharding is not
46// enabled.
47func getSharding() (index, total int, err error) {
48	statusFile := os.Getenv(shardPrefix + shardStatusFileEnv)
49	totalNumStr := os.Getenv(shardPrefix + shardTotalEnv)
50	indexStr := os.Getenv(shardPrefix + shardIndexEnv)
51	if len(totalNumStr) == 0 || len(indexStr) == 0 {
52		return 0, 0, nil
53	}
54
55	totalNum, err := strconv.Atoi(totalNumStr)
56	if err != nil {
57		return 0, 0, fmt.Errorf("$%s is %q, but expected a number\n", shardTotalEnv, totalNumStr)
58	}
59
60	index, err = strconv.Atoi(indexStr)
61	if err != nil {
62		return 0, 0, fmt.Errorf("$%s is %q, but expected a number\n", shardIndexEnv, indexStr)
63	}
64
65	if index < 0 || index >= totalNum {
66		return 0, 0, fmt.Errorf("shard index/total of %d/%d is invalid\n", index, totalNum)
67	}
68
69	if len(statusFile) > 0 {
70		if err := os.WriteFile(statusFile, nil, 0664); err != nil {
71			return 0, 0, err
72		}
73	}
74
75	return index, totalNum, nil
76}
77