diff --git a/.asf.yaml b/.asf.yaml index 21ee08f7a..fabd675fa 100644 --- a/.asf.yaml +++ b/.asf.yaml @@ -38,10 +38,7 @@ github: # strict means "Require branches to be up-to-date before merging". strict: false # contexts are the names of checks that must pass - contexts: - #- computer-ci # always stuck in CI - - check-license-header - - Analyze (java) + contexts: [] required_pull_request_reviews: dismiss_stale_reviews: true require_code_owner_reviews: false diff --git a/.github/workflows/vermeer-ci.yml b/.github/workflows/vermeer-ci.yml index 777ecfcec..bc93ca4cc 100644 --- a/.github/workflows/vermeer-ci.yml +++ b/.github/workflows/vermeer-ci.yml @@ -72,6 +72,14 @@ jobs: - name: Generate embedded assets run: cd asset && go generate + - name: Test deployment credentials and graph initialization failures + run: | + go test -race ./apps/common + go test -race ./apps/master/workers ./apps/master/services -run 'TestWorkerAvailability' + go test -race ./apps/master/bl -run 'Test(WorkerAvailability|HugegraphLoadAuthorization)' + go test -race ./apps/worker -run 'Test(LoadReceivers|RecvVertex|ComputeDispatchCancellation)' + go test -race ./apps/structure -run 'Test(VertexStore|GraphPropagates|MemoryVertex|HugegraphSchema)' + - name: Build run: CGO_ENABLED=0 go build -o vermeer @@ -91,13 +99,14 @@ jobs: uses: actions/checkout@v4 - name: Set up Docker Buildx - uses: docker/setup-buildx-action@v3 + uses: docker/setup-buildx-action@f87e5991a6d7451dcb8d9637bfbc97413f497069 # v4.4.1 - name: Build Docker image (dry-run) - uses: docker/build-push-action@v6 + uses: docker/build-push-action@c3c9e263c25d99ce0380d002d59b67737d91b0dc # v7.4.0 with: context: vermeer push: false + load: true tags: vermeer:ci-check cache-from: type=gha cache-to: type=gha,mode=max diff --git a/vermeer/README.md b/vermeer/README.md index 97cf58d71..b1990d39f 100644 --- a/vermeer/README.md +++ b/vermeer/README.md @@ -241,7 +241,8 @@ auth = none auth_token_factor = 1234 ``` -**Note**: HugeGraph connection details (`pd_peers`, `server`, `graph`) are provided in the graph load API request, not in the configuration file. See [HugeGraph Integration](#hugegraph-integration) section for details. +**Note**: Graph names and legacy HugeGraph connection details are supplied in the graph load request. +Deployment credentials require the trusted PD configuration described below. ### Generating Auth Token @@ -269,6 +270,55 @@ The generated token can be used in two ways: ### Worker Configuration (`worker.ini`) +For a PD deployment using a shared secret, configure both Master and Worker with +`pd_password_file=/path/to/pd-password`, or set `VERMEER_PD_PASSWORD_FILE` to that +file. Its contents must match PD's `auth.secret-key`; protect the file with mode +`0600`. The environment variable takes precedence. A configured unreadable or +empty file fails the PD request instead of using legacy credentials. The secret +is sent as the internal `vermeer` credential and is never added to task parameters. +When no file is configured, the legacy credential remains available for older PD +deployments; it does not authenticate to a modern PD with a different shared secret. + +When either deployment credential file is configured, configure both Master and +Worker with `pd_peers=["pd.internal:8686"]`, or set `VERMEER_PD_PEERS` to that JSON +array (environment takes precedence). These deployment addresses replace task +PD addresses for load and output discovery. Missing or invalid deployment peers +fail before discovery; no credential file is read for PD discovery in that case. +Schema service credentials are used only with Servers discovered through these +PD addresses, and Schema HTTP requests do not follow redirects. Without deployment +credentials or peers, task-provided PD addresses retain their legacy behavior. + +With Master `auth=token`, HugeGraph bulk load and reload tasks require the existing +Vermeer administrator role, including through the native async, sync, batch and +OLTP task APIs. Direct Store scans do not enforce the caller's HugeGraph graph permissions; a Vermeer +space token alone does not authorize a source graph. This restriction applies +with or without deployment credential files, including credentials configured +only on Workers. It is checked before task creation and again before execution +using the recorded task creator and the reload's effective loading type. Existing +non-admin HugeGraph load/reload tasks are rejected when started; administrators +must resubmit them. CSV loads and computation on admitted graphs retain their +existing access rules. + +Master `auth=none` retains its existing trust-all API behavior, including admin +APIs, and provides no caller-level source authorization. Use it only with trusted +callers; use `auth=token` and administrator tokens when source access must be restricted. + +The current PD gRPC transport is plaintext. Keep these connections on a trusted +private network or loopback; deployments crossing untrusted networks must provide +transport encryption. The password file configures request credentials, not TLS. + +For authenticated HugeGraph Schema access, configure the Worker with +`hugegraph_credentials_file=/path/to/hugegraph-credentials.json`, or set +`VERMEER_HUGEGRAPH_CREDENTIALS_FILE` (which takes precedence). The protected file +(mode `0600`) contains `{"username":"schema-reader","password":"..."}` and must +belong to a Server account allowed to read the loaded graphs' Schema. Hubble's +direct adapter never forwards the user's login password. Worker Schema discovery +and reads use these deployment credentials without adding them to persisted task +parameters. A configured missing, unreadable, empty or invalid file fails the +load; it never falls back to task credentials. Without a file, the existing +`load.hugegraph_username` / `load.hugegraph_password` behavior is retained. This +setting does not configure the separate HugeGraph result writer. + ```ini [default] # Worker HTTP listen address @@ -327,8 +377,26 @@ Vermeer exposes a REST API on port `6688` (configurable in `master.ini`). | `/api/v1/compute` | POST | Execute algorithm on loaded graph | | `/api/v1/tasks/{task_id}` | GET | Get task status and results | | `/api/v1/workers` | GET | List connected workers | +| `/api/v1/workers/availability` | GET | Read worker readiness for the authenticated space | | `/ui/` | GET | Web UI dashboard | +Worker availability uses the normal space credential and returns only +`{"available":true|false}` for that authenticated space's allocations or the +common pool. It accepts no caller-selected space and does not allocate workers +or require an administrator token. Global worker registration alone does not +establish readiness for a particular space. +A group preallocated to a graph absent from that space's Master graph metadata +is not counted until the graph appears there; the scheduler can still use the +allocation for that graph's first load. + +If compute dispatch fails after some workers accept it, Master records the original +error and makes a best-effort cancellation request to each accepted worker's original +session, with a five-second deadline per worker. A worker that has not created the +task yet returns `NotFound`, which this rollback retries within the same deadline. +Compute initialization and the first-superstep handoff consume cancellation even +while the worker is idle; active BSP and graph preparation retain their existing +cooperative cancellation behavior. + ### Example: Run PageRank ```bash diff --git a/vermeer/apps/common/http_server.go b/vermeer/apps/common/http_server.go index 5a5c3d28f..a9a0f8502 100644 --- a/vermeer/apps/common/http_server.go +++ b/vermeer/apps/common/http_server.go @@ -33,6 +33,10 @@ type Sentinel struct { srv *http.Server } +func (sen *Sentinel) ServeHTTP(w http.ResponseWriter, r *http.Request) { + sen.engine.ServeHTTP(w, r) +} + func (sen *Sentinel) Init(peer string, debugMode string) { if debugMode == "release" { gin.SetMode(gin.ReleaseMode) diff --git a/vermeer/apps/common/hugegraph_credentials.go b/vermeer/apps/common/hugegraph_credentials.go new file mode 100644 index 000000000..5f5633d2c --- /dev/null +++ b/vermeer/apps/common/hugegraph_credentials.go @@ -0,0 +1,47 @@ +/* +Licensed to the Apache Software Foundation (ASF) under one or more +contributor license agreements. See the NOTICE file distributed with this +work for additional information regarding copyright ownership. The ASF +licenses this file to You under the Apache License, Version 2.0 (the +"License"); you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package common + +import ( + "encoding/json" + "errors" + "fmt" + "os" +) + +// ResolveHugegraphCredentials keeps deployment credentials out of persisted task parameters. +// Without a configured file, legacy callers retain their task-provided credentials. +func ResolveHugegraphCredentials(username, password string) (string, string, error) { + file := deploymentConfig("VERMEER_HUGEGRAPH_CREDENTIALS_FILE", "hugegraph_credentials_file") + if file == "" { + return username, password, nil + } + data, err := os.ReadFile(file) + if err != nil { + return "", "", fmt.Errorf("read HugeGraph credentials file: %w", err) + } + var credentials struct { + Username string `json:"username"` + Password string `json:"password"` + } + if json.Unmarshal(data, &credentials) != nil || credentials.Username == "" || credentials.Password == "" { + // Do not include decoder errors or input bytes in task errors or logs. + return "", "", errors.New("HugeGraph credentials file must contain a nonempty username and password") + } + return credentials.Username, credentials.Password, nil +} diff --git a/vermeer/apps/common/hugegraph_credentials_test.go b/vermeer/apps/common/hugegraph_credentials_test.go new file mode 100644 index 000000000..c5adae484 --- /dev/null +++ b/vermeer/apps/common/hugegraph_credentials_test.go @@ -0,0 +1,97 @@ +/* +Licensed to the Apache Software Foundation (ASF) under one or more +contributor license agreements. See the NOTICE file distributed with this +work for additional information regarding copyright ownership. The ASF +licenses this file to You under the Apache License, Version 2.0 (the +"License"); you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package common + +import ( + "os" + "path/filepath" + "strings" + "testing" +) + +func TestHugegraphCredentialsFilePrecedence(t *testing.T) { + old := GetConfigDefault("hugegraph_credentials_file", "") + t.Cleanup(func() { SetConfig("hugegraph_credentials_file", old) }) + configured := filepath.Join(t.TempDir(), "configured.json") + override := filepath.Join(t.TempDir(), "override.json") + for file, data := range map[string]string{ + configured: `{"username":"config-user","password":"config-secret"}`, + override: `{"username":"environment-user","password":"environment-secret"}`, + } { + if err := os.WriteFile(file, []byte(data), 0600); err != nil { + t.Fatal(err) + } + } + SetConfig("hugegraph_credentials_file", configured) + t.Setenv("VERMEER_HUGEGRAPH_CREDENTIALS_FILE", override) + username, password, err := ResolveHugegraphCredentials("task-user", "task-secret") + if err != nil || username != "environment-user" || password != "environment-secret" { + t.Fatal("environment file must override config and task credentials") + } + t.Setenv("VERMEER_HUGEGRAPH_CREDENTIALS_FILE", "") + username, password, err = ResolveHugegraphCredentials("task-user", "task-secret") + if err != nil || username != "config-user" || password != "config-secret" { + t.Fatal("configuration file must override task credentials") + } + SetConfig("hugegraph_credentials_file", "") + username, password, err = ResolveHugegraphCredentials("task-user", "task-secret") + if err != nil || username != "task-user" || password != "task-secret" { + t.Fatal("unconfigured deployments must retain legacy task credentials") + } + username, password, err = ResolveHugegraphCredentials("", "") + if err != nil || username != "" || password != "" { + t.Fatal("unconfigured anonymous deployments must remain usable") + } +} + +func TestHugegraphCredentialsFileFailsClosed(t *testing.T) { + old := GetConfigDefault("hugegraph_credentials_file", "") + t.Cleanup(func() { SetConfig("hugegraph_credentials_file", old) }) + fallback := filepath.Join(t.TempDir(), "valid-fallback.json") + if err := os.WriteFile(fallback, []byte(`{"username":"config-user","password":"config-password"}`), 0600); err != nil { + t.Fatal(err) + } + SetConfig("hugegraph_credentials_file", fallback) + for _, test := range []struct{ name, data string }{ + {"missing", ""}, {"empty", ""}, {"malformed", `{"username":"secret-user","password":"secret-password"`}, + {"null", `null`}, {"username-missing", `{"password":"secret-password"}`}, + {"username-empty", `{"username":"","password":"secret-password"}`}, + {"password-missing", `{"username":"secret-user"}`}, + {"password-empty", `{"username":"secret-user","password":""}`}, + {"wrong-type", `{"username":"secret-user","password":123}`}, + } { + t.Run(test.name, func(t *testing.T) { + file := filepath.Join(t.TempDir(), "credentials.json") + if test.name != "missing" { + if err := os.WriteFile(file, []byte(test.data), 0600); err != nil { + t.Fatal(err) + } + } + t.Setenv("VERMEER_HUGEGRAPH_CREDENTIALS_FILE", file) + username, password, err := ResolveHugegraphCredentials("task-user", "task-secret") + if err == nil || username != "" || password != "" { + t.Fatal("invalid configured credentials must not fall back to task credentials") + } + for _, secret := range []string{"secret-user", "secret-password", "task-secret"} { + if strings.Contains(err.Error(), secret) { + t.Fatal("credential contents must not appear in errors") + } + } + }) + } +} diff --git a/vermeer/apps/common/hugegraph_tools.go b/vermeer/apps/common/hugegraph_tools.go index c403fa3a4..dc40f991b 100644 --- a/vermeer/apps/common/hugegraph_tools.go +++ b/vermeer/apps/common/hugegraph_tools.go @@ -19,11 +19,13 @@ package common import ( "context" + "encoding/base64" "encoding/json" "errors" "fmt" "io" "net/http" + "os" "strconv" "strings" "sync" @@ -52,9 +54,12 @@ const ( type HgVertexIDType uint16 func FindValidPD(pdIPAddress []string) (string, error) { + pdIPAddress, err := resolvePDPeers(pdIPAddress) + if err != nil { + return "", err + } // retry 3 times to find valid pd address var pdIPAddr string - var err error for i := 0; i < 3; i++ { pdIPAddr, err = findValidPD(pdIPAddress) if err == nil && pdIPAddr != "" { @@ -71,8 +76,12 @@ func findValidPD(pdIPAddress []string) (string, error) { wg := &sync.WaitGroup{} var pdIPAddr string tempCtx, tempCancel := context.WithTimeout(context.Background(), 5*time.Second) + defer tempCancel() pdAuthority := PDAuthority{} - tempCtx = pdAuthority.SetAuthority(tempCtx) + tempCtx, err := pdAuthority.SetAuthority(tempCtx) + if err != nil { + return "", err + } for _, addr := range pdIPAddress { wg.Add(1) go func(ctx context.Context, cancel context.CancelFunc, addr string) { @@ -103,6 +112,9 @@ func findValidPD(pdIPAddress []string) (string, error) { } func FindServerAddr(pdIPAddress string, hgName string, username string, password string) (string, error) { + if err := requireTrustedPD(pdIPAddress); err != nil { + return "", err + } // retry 3 times to find valid server address var serverAddr string var err error @@ -129,7 +141,10 @@ func findServerAddr(pdIPAddress string, hgName string, username string, password } defer pdConn.Close() pdAuthority := PDAuthority{} - ctx = pdAuthority.SetAuthority(ctx) + ctx, err = pdAuthority.SetAuthority(ctx) + if err != nil { + return "", err + } pdClient := pd.NewDiscoveryServiceClient(pdConn) resp, err := pdClient.GetNodes(ctx, &pd.Query{}) if err != nil { @@ -180,12 +195,12 @@ func testServerIsValid(serverAdds []string, hgSpace, hGraph, username, password defer wg.Done() url := fmt.Sprintf("%v/graphspaces/%v/graphs/%v/schema?format=json", addr, hgSpace, hGraph) req, err := http.NewRequest(http.MethodGet, url, nil) - req.SetBasicAuth(username, password) if err != nil { return } - req.WithContext(ctx) - client := &http.Client{Timeout: 5 * time.Second} + req.SetBasicAuth(username, password) + req = req.WithContext(ctx) + client := schemaHTTPClient() resp, err := client.Do(req) if err != nil { return @@ -224,15 +239,15 @@ func GetHugegraphSchema(serverAddr, hgName, username, password string) (map[stri func getHugegraphSchema(url, username, password string) (map[string]any, error) { req, err := http.NewRequest(http.MethodGet, url, nil) - req.SetBasicAuth(username, password) if err != nil { logrus.Errorf("create http request failed, %v", err.Error()) return nil, fmt.Errorf("create http request failed, %v", err.Error()) } + req.SetBasicAuth(username, password) ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - req.WithContext(ctx) - client := &http.Client{Timeout: 5 * time.Second} + req = req.WithContext(ctx) + client := schemaHTTPClient() resp, err := client.Do(req) if err != nil { logrus.Errorf("http request failed, %v", err.Error()) @@ -256,6 +271,15 @@ func getHugegraphSchema(url, username, password string) (map[string]any, error) return propertyFromHg, nil } +func schemaHTTPClient() *http.Client { + return &http.Client{ + Timeout: 5 * time.Second, + CheckRedirect: func(*http.Request, []*http.Request) error { + return http.ErrUseLastResponse + }, + } +} + func SplitHgName(hgName string) (string, string, error) { hgSplit := strings.Split(hgName, "/") if len(hgSplit) != 3 { @@ -320,18 +344,31 @@ type PDAuthority struct { locker sync.RWMutex } -func (pa *PDAuthority) SetAuthority(ctx context.Context) context.Context { +func (pa *PDAuthority) SetAuthority(ctx context.Context) (context.Context, error) { pa.locker.RLock() defer pa.locker.RUnlock() + credential := CredentialValue // Legacy PD compatibility when no file is configured. + passwordFile := deploymentConfig("VERMEER_PD_PASSWORD_FILE", "pd_password_file") + if passwordFile != "" { + password, err := os.ReadFile(passwordFile) + if err != nil { + return ctx, fmt.Errorf("read PD password file: %w", err) + } + secret := strings.TrimSuffix(strings.TrimSuffix(string(password), "\n"), "\r") + if secret == "" { + return ctx, errors.New("PD password file is empty") + } + credential = base64.StdEncoding.EncodeToString([]byte("vermeer:" + secret)) + } md := metadata.New(map[string]string{ - CredentialKey: CredentialValue, + CredentialKey: credential, }) if pa.token != "" { md.Set(TokenKey, pa.token) } // logrus.Infof("send md:%v", md) ctx = metadata.NewOutgoingContext(ctx, md) - return ctx + return ctx, nil } func (pa *PDAuthority) SetToken(md metadata.MD) { diff --git a/vermeer/apps/common/pd_authority_test.go b/vermeer/apps/common/pd_authority_test.go new file mode 100644 index 000000000..c1395d7e8 --- /dev/null +++ b/vermeer/apps/common/pd_authority_test.go @@ -0,0 +1,65 @@ +/* +Licensed to the Apache Software Foundation (ASF) under one or more +contributor license agreements. See the NOTICE file distributed with this +work for additional information regarding copyright ownership. The ASF +licenses this file to You under the Apache License, Version 2.0 (the +"License"); you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package common + +import ( + "context" + "encoding/base64" + "os" + "path/filepath" + "testing" + + "google.golang.org/grpc/metadata" +) + +func TestPDAuthorityPasswordFile(t *testing.T) { + passwordFile := filepath.Join(t.TempDir(), "pd-password") + if err := os.WriteFile(passwordFile, []byte("test:shared-secret\n"), 0600); err != nil { + t.Fatal(err) + } + t.Setenv("VERMEER_PD_PASSWORD_FILE", passwordFile) + authority := &PDAuthority{} + ctx, err := authority.SetAuthority(context.Background()) + if err != nil { + t.Fatal(err) + } + md, _ := metadata.FromOutgoingContext(ctx) + values := md.Get(CredentialKey) + if len(values) != 1 { + t.Fatal("missing PD credential") + } + decoded, err := base64.StdEncoding.DecodeString(values[0]) + if err != nil || string(decoded) != "vermeer:test:shared-secret" { + t.Fatal("PD credential must use the configured shared secret") + } +} + +func TestPDAuthorityRejectsUnavailablePasswordFile(t *testing.T) { + for _, missing := range []bool{true, false} { + file := filepath.Join(t.TempDir(), "password") + if !missing { + if err := os.WriteFile(file, nil, 0600); err != nil { + t.Fatal(err) + } + } + t.Setenv("VERMEER_PD_PASSWORD_FILE", file) + if _, err := (&PDAuthority{}).SetAuthority(context.Background()); err == nil { + t.Fatal("must not fall back to legacy credentials for an invalid configured file") + } + } +} diff --git a/vermeer/apps/common/pd_peers.go b/vermeer/apps/common/pd_peers.go new file mode 100644 index 000000000..3812c3f8b --- /dev/null +++ b/vermeer/apps/common/pd_peers.go @@ -0,0 +1,78 @@ +/* +Licensed to the Apache Software Foundation (ASF) under one or more +contributor license agreements. See the NOTICE file distributed with this +work for additional information regarding copyright ownership. The ASF +licenses this file to You under the Apache License, Version 2.0 (the +"License"); you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package common + +import ( + "encoding/json" + "errors" + "net" + "os" + "slices" + "strconv" + + "github.com/sirupsen/logrus" +) + +func deploymentConfig(environment, key string) string { + if value := os.Getenv(environment); value != "" { + return value + } + value, _ := GetConfigDefault(key, "").(string) + return value +} + +// Deployment discovery takes precedence over task parameters on both Master and Worker. +func resolvePDPeers(taskPeers []string) ([]string, error) { + configured := deploymentConfig("VERMEER_PD_PEERS", "pd_peers") + if configured == "" { + if deploymentConfig("VERMEER_PD_PASSWORD_FILE", "pd_password_file") != "" || + deploymentConfig("VERMEER_HUGEGRAPH_CREDENTIALS_FILE", "hugegraph_credentials_file") != "" { + return nil, errors.New("configure pd_peers when using deployment credentials") + } + return taskPeers, nil + } + var peers []string + if json.Unmarshal([]byte(configured), &peers) != nil || len(peers) == 0 { + return nil, errors.New("pd_peers must be a nonempty JSON array of host:port addresses") + } + for _, peer := range peers { + host, port, err := net.SplitHostPort(peer) + number, portErr := strconv.Atoi(port) + if err != nil || host == "" || portErr != nil || number < 1 || number > 65535 { + return nil, errors.New("pd_peers must contain valid host:port addresses") + } + } + for _, peer := range taskPeers { + if !slices.Contains(peers, peer) { + logrus.Warn("configured pd_peers override task PD peers") + break + } + } + return peers, nil +} + +func requireTrustedPD(peer string) error { + peers, err := resolvePDPeers([]string{peer}) + if err != nil { + return err + } + if !slices.Contains(peers, peer) { + return errors.New("PD address is outside configured pd_peers") + } + return nil +} diff --git a/vermeer/apps/common/pd_peers_test.go b/vermeer/apps/common/pd_peers_test.go new file mode 100644 index 000000000..a79e559c1 --- /dev/null +++ b/vermeer/apps/common/pd_peers_test.go @@ -0,0 +1,172 @@ +/* +Licensed to the Apache Software Foundation (ASF) under one or more +contributor license agreements. See the NOTICE file distributed with this +work for additional information regarding copyright ownership. The ASF +licenses this file to You under the Apache License, Version 2.0 (the +"License"); you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package common + +import ( + "net/http" + "net/http/httptest" + "path/filepath" + "reflect" + "strings" + "sync/atomic" + "testing" + + "github.com/sirupsen/logrus" + logtest "github.com/sirupsen/logrus/hooks/test" +) + +func TestPDPeersOverrideWarning(t *testing.T) { + isolatePDPeerConfig(t) + logger := logrus.StandardLogger() + previousHooks := logger.ReplaceHooks(make(logrus.LevelHooks)) + previousLevel := logger.GetLevel() + logger.SetLevel(logrus.WarnLevel) + hook := logtest.NewLocal(logger) + t.Cleanup(func() { logger.ReplaceHooks(previousHooks); logger.SetLevel(previousLevel) }) + configured := []string{"deployment-a:8686", "deployment-b:8686"} + t.Setenv("VERMEER_PD_PEERS", `["deployment-a:8686","deployment-b:8686"]`) + for _, peers := range [][]string{nil, {configured[1], configured[0]}, {configured[0]}} { + if _, err := resolvePDPeers(peers); err != nil { + t.Fatal(err) + } + } + if err := requireTrustedPD(configured[0]); err != nil { + t.Fatal(err) + } + if len(hook.AllEntries()) != 0 { + t.Fatal("configured task peers and internal trust checks must remain quiet") + } + taskPeers := []string{"sensitive-task-a:8686", "sensitive-task-b:8686"} + peers, err := resolvePDPeers(taskPeers) + if err != nil || !reflect.DeepEqual(peers, configured) { + t.Fatal("warning must preserve deployment precedence") + } + entries := hook.AllEntries() + if len(entries) != 1 || entries[0].Level != logrus.WarnLevel { + t.Fatalf("expected one override warning, got %d entries", len(entries)) + } + entry, err := entries[0].String() + if err != nil { + t.Fatal(err) + } + for _, target := range append(taskPeers, configured...) { + if strings.Contains(entry, target) { + t.Fatal("override warning exposed a target") + } + } +} + +func isolatePDPeerConfig(t *testing.T) { + t.Helper() + for _, key := range []string{"pd_peers", "pd_password_file", "hugegraph_credentials_file"} { + old := GetConfigDefault(key, "") + SetConfig(key, "") + t.Cleanup(func() { SetConfig(key, old) }) + } + for _, key := range []string{"VERMEER_PD_PEERS", "VERMEER_PD_PASSWORD_FILE", "VERMEER_HUGEGRAPH_CREDENTIALS_FILE"} { + t.Setenv(key, "") + } +} + +func TestPDPeersDeploymentPrecedence(t *testing.T) { + isolatePDPeerConfig(t) + taskPeers := []string{"task-pd:8686"} + peers, err := resolvePDPeers(taskPeers) + if err != nil || !reflect.DeepEqual(peers, taskPeers) { + t.Fatal("unconfigured deployments must retain legacy task discovery") + } + SetConfig("pd_peers", `["configured-pd:8686"]`) + peers, err = resolvePDPeers(taskPeers) + if err != nil || !reflect.DeepEqual(peers, []string{"configured-pd:8686"}) { + t.Fatal("deployment peers must replace task peers") + } + t.Setenv("VERMEER_PD_PEERS", `["environment-pd:8686"]`) + peers, err = resolvePDPeers(taskPeers) + if err != nil || !reflect.DeepEqual(peers, []string{"environment-pd:8686"}) { + t.Fatal("environment peers must override configuration and task peers") + } +} + +func TestPDPeersDeploymentCredentialsRequireConfig(t *testing.T) { + for _, key := range []string{"pd_password_file", "hugegraph_credentials_file", + "VERMEER_PD_PASSWORD_FILE", "VERMEER_HUGEGRAPH_CREDENTIALS_FILE"} { + t.Run(key, func(t *testing.T) { + isolatePDPeerConfig(t) + missing := filepath.Join(t.TempDir(), "not-created") + if strings.HasPrefix(key, "VERMEER_") { + t.Setenv(key, missing) + } else { + SetConfig(key, missing) + } + // No endpoint and no credential file exist: validation must precede IO. + if _, err := FindValidPD(nil); err == nil || !strings.Contains(err.Error(), "configure pd_peers") { + t.Fatalf("expected deployment peer validation before discovery, got %v", err) + } + if err := requireTrustedPD("task-pd:8686"); err == nil { + t.Fatal("Server discovery must require deployment PD configuration too") + } + }) + } +} + +func TestPDPeersInvalidDeploymentFailsClosed(t *testing.T) { + isolatePDPeerConfig(t) + SetConfig("pd_peers", `["configured-pd:8686"]`) + for _, value := range []string{`null`, `[]`, `not-json`, `[""]`, `[null]`, + `["task-pd"]`, `["task-pd:0"]`, `["task-pd:65536"]`, `["task-pd:8686",""]`} { + t.Run(value, func(t *testing.T) { + t.Setenv("VERMEER_PD_PEERS", value) + if _, err := FindValidPD([]string{"task-pd:8686"}); err == nil { + t.Fatal("invalid deployment peers must not fall back to config or task peers") + } + }) + } +} + +func TestPDPeersConstrainServerDiscovery(t *testing.T) { + isolatePDPeerConfig(t) + t.Setenv("VERMEER_PD_PEERS", `["configured-pd:8686"]`) + if err := requireTrustedPD("configured-pd:8686"); err != nil { + t.Fatal(err) + } + if _, err := FindServerAddr("task-pd:8686", "DEFAULT/graph/g", "fixture-user", "fixture-password"); err == nil || !strings.Contains(err.Error(), "outside configured pd_peers") { + t.Fatal("Server discovery must reject an unconfigured PD before dialing") + } +} + +func TestHugegraphSchemaDoesNotFollowRedirects(t *testing.T) { + var redirected atomic.Int32 + destination := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + redirected.Add(1) + _, _ = w.Write([]byte(`{}`)) + })) + defer destination.Close() + source := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, destination.URL, http.StatusFound) + })) + defer source.Close() + if addr := testServerIsValid([]string{source.URL}, "DEFAULT", "graph", "fixture-user", "fixture-password"); addr != "" { + t.Fatal("a redirect must not count as a discovered Schema Server") + } + if _, err := getHugegraphSchema(source.URL, "fixture-user", "fixture-password"); err == nil { + t.Fatal("a Schema redirect must be reported as an error") + } + if redirected.Load() != 0 { + t.Fatal("Schema discovery and reads must stay at the original Server") + } +} diff --git a/vermeer/apps/compute/task.go b/vermeer/apps/compute/task.go index b73cd80ef..486b7a933 100644 --- a/vermeer/apps/compute/task.go +++ b/vermeer/apps/compute/task.go @@ -46,6 +46,9 @@ type ComputerTask struct { SettingOutEdges bool SendCount map[string]*int32 RecvCount map[string]*int32 + // WaitingForStep is protected by Locker during the initialization handoff. + WaitingForStep bool + CancelOnce sync.Once } func (ct *ComputerTask) FreeMemory() { diff --git a/vermeer/apps/graphio/hugegraph.go b/vermeer/apps/graphio/hugegraph.go index f204324e3..fb01fdf81 100644 --- a/vermeer/apps/graphio/hugegraph.go +++ b/vermeer/apps/graphio/hugegraph.go @@ -78,7 +78,11 @@ func (a *HugegraphMaker) MakeTasks(params map[string]string, taskID int32) ([]Lo return nil, err } pdAuthority := common.PDAuthority{} - ctx = pdAuthority.SetAuthority(ctx) + defer pdConn.Close() + ctx, err = pdAuthority.SetAuthority(ctx) + if err != nil { + return nil, err + } var md metadata.MD pdClient := pd.NewPDClient(pdConn) partitionsResp, err := pdClient.QueryPartitions(ctx, @@ -110,7 +114,10 @@ func (a *HugegraphMaker) MakeTasks(params map[string]string, taskID int32) ([]Lo } } var leaderStoreID uint64 - ctx = pdAuthority.SetAuthority(ctx) + ctx, err = pdAuthority.SetAuthority(ctx) + if err != nil { + return nil, err + } var md metadata.MD shardGroup, err := pdClient.GetShardGroup(ctx, &pd.GetShardGroupRequest{GroupId: partition.Id}, grpc.Header(&md)) if err != nil { @@ -129,7 +136,10 @@ func (a *HugegraphMaker) MakeTasks(params map[string]string, taskID int32) ([]Lo //重复leader不再执行获取地址 if _, ok := leaderStore[leaderStoreID]; !ok { //获得新的leader的地址并写入map - ctx = pdAuthority.SetAuthority(ctx) + ctx, err = pdAuthority.SetAuthority(ctx) + if err != nil { + return nil, err + } var md metadata.MD storeResp, err := pdClient.GetStore(ctx, &pd.GetStoreRequest{StoreId: leaderStoreID}, grpc.Header(&md)) if err != nil { @@ -195,11 +205,6 @@ func (a *HugegraphMaker) MakeTasks(params map[string]string, taskID int32) ([]Lo loadParts[i].Params["load.hugestore_batch_timeout"] = params["load.hugestore_batch_timeout"] loadParts[i].Params["load.hugestore_batchsize"] = params["load.hugestore_batchsize"] } - err = pdConn.Close() - if err != nil { - logrus.Errorf("hugegraph-pd close err:%v", err) - return nil, err - } return loadParts, nil } diff --git a/vermeer/apps/graphio/load_graph.go b/vermeer/apps/graphio/load_graph.go index b09235f21..b52ee21ff 100644 --- a/vermeer/apps/graphio/load_graph.go +++ b/vermeer/apps/graphio/load_graph.go @@ -172,6 +172,10 @@ type LoadGraphTask struct { SendCount map[string]*int32 RecvCount map[string]*int32 Locker *sync.Mutex + // InitializationFailed is immutable once initialization releases LoadWg. + // Later-phase receivers keep a defensive guard; current load-step ordering + // prevents scatter/edge messages after initialization failure. + InitializationFailed bool } func (lt *LoadGraphTask) GetPartition(partIdx int32) LoadPartition { @@ -452,6 +456,7 @@ func (lw *loadGraphWorker) InstallTask(taskInfo *structure.TaskInfo) *LoadGraphT lw.locker.Lock() defer lw.locker.Unlock() if _, ok := lw.loadTasks[taskInfo.ID]; ok { + lw.loadTasks[taskInfo.ID].LoadWg.Add(1) return lw.loadTasks[taskInfo.ID] } loadTask := LoadGraphTask{ @@ -468,6 +473,9 @@ func (lw *loadGraphWorker) InstallTask(taskInfo *structure.TaskInfo) *LoadGraphT loadTask.SendCount[worker.Name] = new(int32) loadTask.RecvCount[worker.Name] = new(int32) } + // Receivers can find this task as soon as it is published. Hold the + // initialization barrier before making it visible to them. + loadTask.LoadWg.Add(1) lw.loadTasks[taskInfo.ID] = &loadTask // TODO: why not a pointer? return &loadTask diff --git a/vermeer/apps/master/bl/compute_dispatch_test.go b/vermeer/apps/master/bl/compute_dispatch_test.go new file mode 100644 index 000000000..6029839b5 --- /dev/null +++ b/vermeer/apps/master/bl/compute_dispatch_test.go @@ -0,0 +1,190 @@ +/* +Licensed to the Apache Software Foundation (ASF) under one or more +contributor license agreements. See the NOTICE file distributed with this +work for additional information regarding copyright ownership. The ASF +licenses this file to You under the Apache License, Version 2.0 (the +"License"); you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package bl + +import ( + "context" + "errors" + "io" + "testing" + "time" + "vermeer/apps/common" + "vermeer/apps/compute" + "vermeer/apps/master/workers" + pb "vermeer/apps/protos" + "vermeer/apps/structure" + + "github.com/prometheus/client_golang/prometheus" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +type computeDispatchStream struct { + grpc.ServerStream + requests int + sendErr error +} + +func (s *computeDispatchStream) Send(*pb.ComputeTaskResp) error { + s.requests++ + return s.sendErr +} +func (s *computeDispatchStream) Recv() (*pb.ComputeTaskReq, error) { return nil, io.EOF } + +type computeCancelClient struct { + pb.WorkerClient + task *structure.TaskInfo + calls int + notFound bool + cancelErr error + invalid bool +} + +func (c *computeCancelClient) ControlTask(ctx context.Context, req *pb.ControlTaskReq, + _ ...grpc.CallOption) (*pb.ControlTaskResp, error) { + c.calls++ + if req.TaskID != c.task.ID || req.Action != structure.ActionCancelTask { + c.invalid = true + return nil, errors.New("wrong cancellation scope") + } + if c.task.State != structure.TaskStateError { + c.invalid = true + return nil, errors.New("task was not failed before cancellation") + } + deadline, ok := ctx.Deadline() + if !ok || time.Until(deadline) > 5*time.Second { + c.invalid = true + return nil, errors.New("cancellation has no bounded deadline") + } + if c.notFound && c.calls == 1 { + return nil, status.Error(codes.NotFound, "not initialized") + } + return &pb.ControlTaskResp{}, c.cancelErr +} + +// Reuse the registration regression's isolated stores rather than reopen Pebble fixtures. +func testPartialComputeDispatch(t *testing.T) { + previousGraphs := common.PrometheusMetrics.GraphCnt + common.PrometheusMetrics.GraphCnt = prometheus.NewGaugeVec( + prometheus.GaugeOpts{Name: "test_compute_dispatch_graphs"}, nil) + t.Cleanup(func() { common.PrometheusMetrics.GraphCnt = previousGraphs }) + if err := graphMgr.InitStore(); err != nil { + t.Fatal(err) + } + algorithmMgr.Init() + algorithmMgr.Register(new(compute.AlgorithmMakerBase), "test") + for _, test := range []struct { + name string + action pb.ComputeAction + missingWorker, missingStream, sendFailure, cancelFailure, lateInit bool + }{ + {name: "missing_worker", missingWorker: true}, + {name: "missing_stream", missingStream: true}, + {name: "send_failure", sendFailure: true}, + {name: "cleanup_failure", missingStream: true, cancelFailure: true}, + {name: "late_initialization", missingStream: true, lateInit: true}, + {name: "success"}, + {name: "out_edges", action: pb.ComputeAction_SettingOutEdges, sendFailure: true}, + {name: "out_degree", action: pb.ComputeAction_SettingOutDegree, sendFailure: true}, + } { + t.Run(test.name, func(t *testing.T) { + ServerMgr.Init() + defer ServerMgr.Init() + task, err := taskMgr.CreateTask("compute-dispatch", structure.TaskTypeCompute, 0) + if err != nil { + t.Fatal(err) + } + task.GraphName = "graph" + graph := graphMgr.CreateGraph(task.SpaceName, task.GraphName) + clients := make([]*computeCancelClient, 3) + streams := make([]*computeDispatchStream, 3) + names := []string{"compute-first", "compute-second", "compute-third"} + dispatchErr := errors.New("compute dispatch failed") + for i, name := range names { + graph.Workers = append(graph.Workers, &structure.GraphWorker{Name: name}) + task.Workers = append(task.Workers, &structure.TaskWorker{Name: name}) + clients[i] = &computeCancelClient{task: task} + streams[i] = new(computeDispatchStream) + workerMgr.RemoveWorker(name) + if i == 2 && test.missingWorker { + continue + } + _, err = workerMgr.AddWorker(&workers.WorkerClient{Name: name, Session: clients[i]}) + if err != nil { + t.Fatal(err) + } + if i != 2 || !test.missingStream { + ServerMgr.computeSevers[name] = &ComputeTaskServer{streamServer: streams[i]} + } + } + if test.sendFailure { + streams[2].sendErr = dispatchErr + } + if test.cancelFailure { + clients[0].cancelErr = errors.New("cancel failed") + } + clients[0].notFound = test.lateInit + ct := &compute.ComputerTask{Task: task, Algorithm: "base"} + err = StartComputeTask(graph, ct, test.action) + failed := test.missingWorker || test.missingStream || test.sendFailure + if (err != nil) != failed { + t.Fatalf("dispatch err=%v, want failure=%v", err, failed) + } + if test.sendFailure && !errors.Is(err, dispatchErr) { + t.Fatalf("replaced original error: %v", err) + } + if failed && (task.State != structure.TaskStateError || task.ErrMessage != err.Error()) { + t.Errorf("partial dispatch left master state=%s, error=%q; want Error and original dispatch error", task.State, task.ErrMessage) + } else if failed { + saved, readErr := taskMgr.ReadTask(task.ID) + if readErr != nil { + t.Fatal(readErr) + } + if task.State != structure.TaskStateError || saved.State != task.State || + task.ErrMessage != err.Error() || saved.ErrMessage != err.Error() { + t.Fatal("dispatch failure must persist the original error") + } + } + for i, c := range clients { + if c.invalid { + t.Errorf("worker %d received an invalid cancellation", i) + } + want := 0 + if failed && i < 2 { + want = 1 + } + if test.lateInit && i == 0 { + want = 2 + } + if c.calls != want { + t.Errorf("worker %d cancellation calls=%d, want %d", i, c.calls, want) + } + } + for i, s := range streams { + want := 1 + if i == 2 && (test.missingWorker || test.missingStream) { + want = 0 + } + if s.requests != want { + t.Errorf("worker %d sends=%d want=%d", i, s.requests, want) + } + } + }) + } +} diff --git a/vermeer/apps/master/bl/grpc_handlers.go b/vermeer/apps/master/bl/grpc_handlers.go index 2734c063f..4b5378c73 100644 --- a/vermeer/apps/master/bl/grpc_handlers.go +++ b/vermeer/apps/master/bl/grpc_handlers.go @@ -136,17 +136,18 @@ func (h *ServerHandler) SayHelloMaster(ctx context.Context, req *pb.HelloMasterR grpc.MaxCallSendMsgSize(1*1024*1024*1024), grpc.MaxCallRecvMsgSize(1*1024*1024*1024)) - workerClient.Connection, err = grpc.Dial( + connection, err := grpc.Dial( ip+":"+port, grpc.WithTransportCredentials(insecure.NewCredentials()), dialOptions, grpc.WithBlock(), grpc.WithIdleTimeout(0)) if err != nil { logrus.Errorf("connect worker error: %s", err) return &resp, err } + workerClient.Session = pb.NewWorkerClient(connection) + workerMgr.SetWorkerConnection(workerClient, connection) spaces := structure.SpaceManager.GetAllSpace() for _, space := range spaces { resp.Spaces = append(resp.Spaces, space.Name()) } - workerClient.Session = pb.NewWorkerClient(workerClient.Connection) initWorkerClient(workerClient) @@ -235,11 +236,9 @@ func (h *ServerHandler) SuperStep(req *pb.SuperStepReq, stream pb.Master_SuperSt if err := ServerMgr.PutSuperStepServer(name, stream); err != nil { return err } - - for i := 0; i < 100000000; i++ { - time.Sleep(10 * time.Second) - } - //WorkerMgr.GetWorker(name).SuperStepServer.RecvHandler(name) + worker := workerMgr.GetWorker(name) + defer workerMgr.ClearTaskStream(worker, workers.SuperStepStream, stream.Context()) + <-stream.Context().Done() return nil } @@ -383,6 +382,7 @@ func (h *ServerHandler) UploadTPResult(ctx context.Context, in *pb.UploadTPResul type LoadGraphServer struct { streamServer pb.Master_LoadGraphTaskServer + sendMutex sync.Mutex } func (s *LoadGraphServer) SetServer(stream pb.Master_LoadGraphTaskServer) { @@ -435,6 +435,9 @@ func (s *LoadGraphServer) AsyncLoad( spaceName string, workers []string, params map[string]string) error { + // Dispatch, phase changes and error cleanup share this stream. + s.sendMutex.Lock() + defer s.sendMutex.Unlock() err := s.streamServer.Send(&pb.LoadGraphTaskResp{ Base: &pb.BaseResponse{}, LoadType: loadType, diff --git a/vermeer/apps/master/bl/load_authorization_test.go b/vermeer/apps/master/bl/load_authorization_test.go new file mode 100644 index 000000000..47f584b44 --- /dev/null +++ b/vermeer/apps/master/bl/load_authorization_test.go @@ -0,0 +1,165 @@ +/* +Licensed to the Apache Software Foundation (ASF) under one or more +contributor license agreements. See the NOTICE file distributed with this +work for additional information regarding copyright ownership. The ASF +licenses this file to You under the Apache License, Version 2.0 (the +"License"); you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package bl + +import ( + "testing" + "vermeer/apps/common" + "vermeer/apps/structure" + + "github.com/prometheus/client_golang/prometheus" +) + +func isLoadPermissionError(err error) bool { + return err != nil && err.Error() == "admin permission is required to load from HugeGraph" +} + +func setLoadAuthorizationMode(t *testing.T, mode string) { + t.Helper() + previous := common.GetConfigDefault("auth", "none") + common.SetConfig("auth", mode) + t.Cleanup(func() { common.SetConfig("auth", previous) }) +} + +func TestHugegraphLoadAuthorization(t *testing.T) { + setLoadAuthorizationMode(t, "token") + previousGraphs := common.PrometheusMetrics.GraphCnt + common.PrometheusMetrics.GraphCnt = prometheus.NewGaugeVec( + prometheus.GaugeOpts{Name: "test_load_authorization_graphs"}, nil) + graphMgr.Init() + taskMgr.Init() + t.Cleanup(func() { + graphMgr.Init() + taskMgr.Init() + common.PrometheusMetrics.GraphCnt = previousGraphs + }) + graph := graphMgr.CreateGraph("own-space", "alias") + graph.Params = map[string]string{"load.type": "hugegraph", "load.hugegraph_name": "source/graph/g"} + if err := graphMgr.AddGraph(graph); err != nil { + t.Fatal(err) + } + for _, test := range []struct { + name, user, taskType string + params map[string]string + denied bool + }{ + {"user same-space source", "reader", "load", map[string]string{ + "load.type": "hugegraph", "load.hugegraph_name": "own-space/graph/g"}, true}, + {"user other-space source", "reader", "load", map[string]string{ + "load.type": "hugegraph", "load.hugegraph_name": "other-space/graph/g"}, true}, + {"caller params cannot set identity", "reader", "load", map[string]string{ + "load.type": "hugegraph", "create_user": "admin", "user": "admin"}, true}, + {"default user", structure.DefaultUserName, "load", map[string]string{"load.type": "hugegraph"}, true}, + {"admin load", "admin", "load", map[string]string{"load.type": "hugegraph"}, false}, + {"user explicit reload", "reader", "reload", map[string]string{"load.type": "hugegraph"}, true}, + {"user inherited reload", "reader", "reload", nil, true}, + {"admin inherited reload", "admin", "reload", nil, false}, + {"user CSV load", "reader", "load", map[string]string{"load.type": "csv"}, false}, + {"user CSV reload override", "reader", "reload", map[string]string{"load.type": "csv"}, false}, + {"user compute", "reader", "compute", map[string]string{"load.type": "hugegraph"}, false}, + } { + t.Run(test.name, func(t *testing.T) { + cred, err := structure.NewCredential(test.user, "own-space") + if err != nil { + t.Fatal(err) + } + creator, err := NewTaskCreator(cred, test.taskType) + if err != nil { + t.Fatal(err) + } + // This is the real creation preview used before duplicate checks, graph changes or persistence. + task, err := creator.NewTaskInfo("alias", test.params, -1) + if test.denied { + if !isLoadPermissionError(err) || task != nil { + t.Fatalf("unauthorized creation returned task=%v, err=%v", task, err) + } + // Both ordinary and batch admission must reject before graph/task side effects. + for _, checked := range []bool{IsCheck, NoCheck} { + if _, err := (&TaskBl{Cred: cred}).CreateTaskInfo( + "alias", test.taskType, test.params, checked); !isLoadPermissionError(err) { + t.Fatalf("checked=%v: unexpected admission error %v", checked, err) + } + } + } else if err != nil || task.CreateUser != test.user || task.SpaceName != "own-space" { + t.Fatalf("allowed creation lost authenticated identity: task=%v, err=%v", task, err) + } + + // Restored and cron-copied tasks must be checked without relying on a live HTTP credential. + restored := &structure.TaskInfo{CreateUser: test.user, SpaceName: "own-space", + GraphName: "alias", Type: test.taskType, Params: test.params} + starter, err := NewTaskStarter(restored, "$") + if test.denied { + if !isLoadPermissionError(err) || starter != nil { + t.Fatalf("unauthorized execution returned starter=%v, err=%v", starter, err) + } + } else if err != nil || starter == nil { + t.Fatalf("allowed execution rejected: %v", err) + } + }) + } + if len(taskMgr.GetAllTasks(100)) != 0 || graphMgr.GetGraphByName("own-space", "alias") != graph { + t.Fatal("rejected requests changed task or graph state") + } +} + +func TestHugegraphLoadAuthorizationRechecksReload(t *testing.T) { + setLoadAuthorizationMode(t, "token") + graphMgr.Init() + t.Cleanup(graphMgr.Init) + previousGraphs := common.PrometheusMetrics.GraphCnt + common.PrometheusMetrics.GraphCnt = prometheus.NewGaugeVec( + prometheus.GaugeOpts{Name: "test_reload_authorization_graphs"}, nil) + t.Cleanup(func() { common.PrometheusMetrics.GraphCnt = previousGraphs }) + graph := graphMgr.CreateGraph("own-space", "alias") + graph.Params = map[string]string{"load.type": "csv"} + if err := graphMgr.AddGraph(graph); err != nil { + t.Fatal(err) + } + cred, _ := structure.NewCredential("reader", "own-space") + creator, _ := NewTaskCreator(cred, structure.TaskTypeReload) + task, err := creator.NewTaskInfo("alias", nil, -1) + if err != nil { + t.Fatal(err) + } + // A queued reload inherits the graph parameters at execution time. + graph.Params = map[string]string{"load.type": "hugegraph"} + if _, err := NewTaskStarter(task, "$"); !isLoadPermissionError(err) { + t.Fatalf("queued reload inherited a privileged source without rejection: %v", err) + } +} + +func TestHugegraphLoadAuthorizationWithoutTokenAuth(t *testing.T) { + setLoadAuthorizationMode(t, "none") + taskMgr.Init() + t.Cleanup(taskMgr.Init) + creator, err := NewTaskCreator(structure.DefaultCredential(), structure.TaskTypeLoad) + if err != nil { + t.Fatal(err) + } + task, err := creator.NewTaskInfo("graph", map[string]string{"load.type": "hugegraph"}, -1) + if err != nil { + t.Fatalf("no-auth deployments must retain their existing trusted-caller behavior: %v", err) + } + if _, err := NewTaskStarter(task, "$"); err != nil { + t.Fatal(err) + } + common.SetConfig("auth", "token") + if _, err := NewTaskStarter(task, "$"); !isLoadPermissionError(err) { + t.Fatalf("enabling token auth must reject an old non-admin load: %v", err) + } +} diff --git a/vermeer/apps/master/bl/server_manager.go b/vermeer/apps/master/bl/server_manager.go index 6f3467439..e9856c51f 100644 --- a/vermeer/apps/master/bl/server_manager.go +++ b/vermeer/apps/master/bl/server_manager.go @@ -20,6 +20,7 @@ package bl import ( "errors" "sync" + "vermeer/apps/master/workers" pb "vermeer/apps/protos" ) @@ -42,11 +43,21 @@ func (s *ServerManager) StartingLoadServer(workerName string, stream pb.Master_L } loadServer := &LoadGraphServer{streamServer: stream} + s.locker.Lock() s.loadSevers[workerName] = loadServer + worker := workerMgr.GetWorker(workerName) + workerMgr.SetTaskStream(worker, workers.LoadTaskStream, stream.Context()) + s.locker.Unlock() + defer func() { + workerMgr.ClearTaskStream(worker, workers.LoadTaskStream, stream.Context()) + s.locker.Lock() + defer s.locker.Unlock() + if s.loadSevers[workerName] == loadServer { + delete(s.loadSevers, workerName) + } + }() loadServer.RecvHandler(workerName) - delete(s.loadSevers, workerName) - return nil } @@ -56,11 +67,21 @@ func (s *ServerManager) StartingComputeServer(workerName string, stream pb.Maste } computeServer := &ComputeTaskServer{streamServer: stream} + s.locker.Lock() s.computeSevers[workerName] = computeServer + worker := workerMgr.GetWorker(workerName) + workerMgr.SetTaskStream(worker, workers.ComputeTaskStream, stream.Context()) + s.locker.Unlock() + defer func() { + workerMgr.ClearTaskStream(worker, workers.ComputeTaskStream, stream.Context()) + s.locker.Lock() + defer s.locker.Unlock() + if s.computeSevers[workerName] == computeServer { + delete(s.computeSevers, workerName) + } + }() computeServer.RecvHandler(workerName) - delete(s.computeSevers, workerName) - return nil } func (s *ServerManager) PutSuperStepServer(workerName string, stream pb.Master_SuperStepServer) error { @@ -74,15 +95,22 @@ func (s *ServerManager) PutSuperStepServer(workerName string, stream pb.Master_S delete(s.superStepSevers, workerName) superStepServer := &SuperStepServer{streamServer: stream} s.superStepSevers[workerName] = superStepServer + // Keep dispatch installation and readiness publication in the same order. + worker := workerMgr.GetWorker(workerName) + workerMgr.SetTaskStream(worker, workers.SuperStepStream, stream.Context()) return nil } func (s *ServerManager) LoadServer(workerName string) *LoadGraphServer { + s.locker.Lock() + defer s.locker.Unlock() // TODO: handle nil return s.loadSevers[workerName] } func (s *ServerManager) ComputeServer(workerName string) *ComputeTaskServer { + s.locker.Lock() + defer s.locker.Unlock() // TODO: handle nil return s.computeSevers[workerName] } diff --git a/vermeer/apps/master/bl/task_bl.go b/vermeer/apps/master/bl/task_bl.go index 1c7e43c9a..57237f87f 100644 --- a/vermeer/apps/master/bl/task_bl.go +++ b/vermeer/apps/master/bl/task_bl.go @@ -335,7 +335,23 @@ func doLoadingTask(graph *structure.VermeerGraph, task *structure.TaskInfo, work return nil } -func sendLoadingReq(task *structure.TaskInfo, workerClients []*WorkerClient, apply Function[*WorkerClient, *LoadingTaskReq]) error { +func sendLoadingReq(task *structure.TaskInfo, workerClients []*WorkerClient, + apply Function[*WorkerClient, *LoadingTaskReq]) (err error) { + dispatched := make([]*LoadGraphServer, 0, len(workerClients)) + defer func() { + if err == nil { + return + } + taskMgr.SetError(task, err.Error()) + // Stop earlier dispatches using their original streams. + // Cleanup is best effort and must not replace the original dispatch error. + for _, server := range dispatched { + if stopErr := server.AsyncLoad(task.ID, pb.LoadStep_Error, "", task.GraphName, + task.SpaceName, nil, nil); stopErr != nil { + logrus.Errorf("failed to stop partially dispatched load task %d: %v", task.ID, stopErr) + } + } + }() for _, wc := range workerClients { taskWorker := structure.TaskWorker{ @@ -344,12 +360,14 @@ func sendLoadingReq(task *structure.TaskInfo, workerClients []*WorkerClient, app } task.Workers = append(task.Workers, &taskWorker) //err := wc.loadServer.SendLoadReq(apply(wc)) - err := ServerMgr.LoadServer(wc.Name).SendLoadReq(apply(wc)) - if err != nil { - //task.SetState(structure.TaskStateError) - taskMgr.SetError(task, err.Error()) + server := ServerMgr.LoadServer(wc.Name) + if server == nil { + return fmt.Errorf("worker %s load task stream is unavailable", wc.Name) + } + if err := server.SendLoadReq(apply(wc)); err != nil { return err } + dispatched = append(dispatched, server) } return nil diff --git a/vermeer/apps/master/bl/task_creator.go b/vermeer/apps/master/bl/task_creator.go index a7353c85b..dd49b1f4b 100644 --- a/vermeer/apps/master/bl/task_creator.go +++ b/vermeer/apps/master/bl/task_creator.go @@ -18,9 +18,13 @@ under the License. package bl import ( + "errors" "fmt" + "vermeer/apps/common" + "vermeer/apps/graphio" . "vermeer/apps/master/graphs" "vermeer/apps/master/tasks" + "vermeer/apps/options" "vermeer/apps/structure" "github.com/sirupsen/logrus" @@ -32,6 +36,32 @@ const ( NoCheck = false ) +var errHugegraphLoadPermission = errors.New("admin permission is required to load from HugeGraph") + +// Store bulk scans do not enforce the caller's HugeGraph permissions. With token +// authentication, only administrators may initiate them, regardless of where credentials are configured. +func checkHugegraphLoadPermission(user, space, graphName, taskType string, params map[string]string) error { + if common.GetConfigDefault("auth", "none").(string) != "token" { + return nil // The existing no-auth mode trusts every API caller, including admin APIs. + } + if taskType != structure.TaskTypeLoad && taskType != structure.TaskTypeReload { + return nil + } + loadType, specified := params["load.type"] + if taskType == structure.TaskTypeReload && !specified { + if graph := graphMgr.GetGraphByName(space, graphName); graph != nil { + loadType = options.GetString(graph.Params, "load.type") + } + } + if loadType == graphio.LoadTypeHugegraph { + cred, err := structure.NewCredential(user, space) + if err != nil || !cred.IsAdmin() { + return errHugegraphLoadPermission + } + } + return nil +} + // TaskCreator Create a new TaskInfo. type TaskCreator interface { // CreateTaskInfo creates a new TaskInfo from the given task description and adds it to the TaskManager. @@ -87,6 +117,9 @@ func (bc *baseCreator) createTaskInfo(graphName string, params map[string]string } func (bc *baseCreator) NewTaskInfo(graphName string, params map[string]string, taskId int32) (*structure.TaskInfo, error) { + if err := checkHugegraphLoadPermission(bc.cred.User(), bc.cred.Space(), graphName, bc.taskType, params); err != nil { + return nil, err + } task, err := taskMgr.CreateTask(bc.cred.Space(), bc.taskType, taskId) if err != nil { return nil, err diff --git a/vermeer/apps/master/bl/task_starter.go b/vermeer/apps/master/bl/task_starter.go index 5549baeeb..3fb3d4650 100644 --- a/vermeer/apps/master/bl/task_starter.go +++ b/vermeer/apps/master/bl/task_starter.go @@ -18,14 +18,19 @@ under the License. package bl import ( + "context" "fmt" + "time" "vermeer/apps/compute" . "vermeer/apps/master/graphs" "vermeer/apps/master/tasks" + "vermeer/apps/master/workers" pb "vermeer/apps/protos" "vermeer/apps/structure" "github.com/sirupsen/logrus" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" ) // TaskStarter The starter of tasks @@ -38,6 +43,11 @@ func NewTaskStarter(taskInfo *structure.TaskInfo, workerGroup string) (TaskStart if taskInfo == nil { return nil, fmt.Errorf("the argument `taskInfo` should not be nil") } + // Recheck persisted/queued tasks against the current reload source before any graph mutation or scan. + if err := checkHugegraphLoadPermission(taskInfo.CreateUser, taskInfo.SpaceName, + taskInfo.GraphName, taskInfo.Type, taskInfo.Params); err != nil { + return nil, err + } base := baseStarter{task: taskInfo, workerGroup: workerGroup} switch taskInfo.Type { case structure.TaskTypeLoad: @@ -272,7 +282,7 @@ func (cts *computeTaskStarter) StartTask() error { return nil } -func StartComputeTask(graph *structure.VermeerGraph, computeTask *compute.ComputerTask, action pb.ComputeAction) error { +func StartComputeTask(graph *structure.VermeerGraph, computeTask *compute.ComputerTask, action pb.ComputeAction) (err error) { var workerState structure.TaskState switch action { case pb.ComputeAction_Compute: @@ -300,6 +310,26 @@ func StartComputeTask(graph *structure.VermeerGraph, computeTask *compute.Comput workerState = structure.TaskStateSettingOutDegree } + dispatched := make([]*workers.WorkerClient, 0, len(graph.Workers)) + defer func() { + if err == nil { + return + } + taskMgr.SetError(computeTask.Task, err.Error()) + for _, worker := range dispatched { + if worker.Session == nil { + logrus.Errorf("cannot cancel partially dispatched compute task %d: worker %s session is unavailable", + computeTask.Task.ID, worker.Name) + continue + } + cancelErr := cancelComputeTask(worker.Session, computeTask.Task.ID) + if cancelErr != nil { + logrus.Errorf("failed to cancel partially dispatched compute task %d on worker %s: %v", + computeTask.Task.ID, worker.Name, cancelErr) + } + } + }() + for _, w := range graph.Workers { wc := workerMgr.GetWorker(w.Name) if wc == nil { @@ -309,7 +339,11 @@ func StartComputeTask(graph *structure.VermeerGraph, computeTask *compute.Comput } computeTask.Task.SetWorkerState(w.Name, workerState) //err := wc.computeServer.AsyncCompute( - err := ServerMgr.ComputeServer(wc.Name).AsyncCompute( + server := ServerMgr.ComputeServer(wc.Name) + if server == nil { + return fmt.Errorf("worker %s compute task stream is unavailable", wc.Name) + } + err := server.AsyncCompute( computeTask.Algorithm, computeTask.Task.GraphName, computeTask.Task.SpaceName, @@ -320,6 +354,25 @@ func StartComputeTask(graph *structure.VermeerGraph, computeTask *compute.Comput if err != nil { return err } + dispatched = append(dispatched, wc) } return nil } + +func cancelComputeTask(session pb.WorkerClient, taskID int32) error { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + req := &pb.ControlTaskReq{TaskID: taskID, Action: structure.ActionCancelTask} + for { + _, err := session.ControlTask(ctx, req) + if status.Code(err) != codes.NotFound { + return err + } + // A successful stream send can precede task creation on the worker. + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(50 * time.Millisecond): + } + } +} diff --git a/vermeer/apps/master/bl/worker_availability_test.go b/vermeer/apps/master/bl/worker_availability_test.go new file mode 100644 index 000000000..7a46337b2 --- /dev/null +++ b/vermeer/apps/master/bl/worker_availability_test.go @@ -0,0 +1,399 @@ +/* +Licensed to the Apache Software Foundation (ASF) under one or more +contributor license agreements. See the NOTICE file distributed with this +work for additional information regarding copyright ownership. The ASF +licenses this file to You under the Apache License, Version 2.0 (the +"License"); you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +License for the specific language governing permissions and limitations +under the License. +*/ + +package bl + +import ( + "context" + "errors" + "fmt" + "io" + "net" + "sync/atomic" + "testing" + "time" + "vermeer/apps/common" + "vermeer/apps/master/schedules" + "vermeer/apps/master/workers" + pb "vermeer/apps/protos" + "vermeer/apps/structure" + "vermeer/apps/version" + + "github.com/prometheus/client_golang/prometheus" + "google.golang.org/grpc" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/peer" +) + +type registrationWindowWorker struct { + pb.UnimplementedWorkerServer + started chan struct{} + release chan struct{} +} + +type dispatchTestStream struct { + grpc.ServerStream + requests []*pb.LoadGraphTaskResp + sendErr error + stopErr error +} + +func (s *dispatchTestStream) Send(req *pb.LoadGraphTaskResp) error { + s.requests = append(s.requests, req) + if req.Step == pb.LoadStep_Error { + return s.stopErr + } + return s.sendErr +} + +func (s *dispatchTestStream) Recv() (*pb.LoadGraphTaskReq, error) { + return nil, io.EOF +} + +type serializedLoadTestStream struct { + dispatchTestStream + entered chan struct{} + release chan struct{} + second chan struct{} + active atomic.Int32 + peak atomic.Int32 + calls atomic.Int32 +} + +type publicationTestStream struct { + grpc.ServerStream + ctx context.Context + onRead func(int32) + reads atomic.Int32 +} + +func (s *publicationTestStream) Context() context.Context { + if s.onRead != nil { + s.onRead(s.reads.Add(1)) + } + return s.ctx +} + +func (s *publicationTestStream) Send(*pb.SuperStepResp) error { return nil } + +// This checks registration ordering with controlled handlers, not a live reconnect. +func testSuperStepPublication(t *testing.T, worker *workers.WorkerClient) { + ServerMgr.Init() + workerMgr.SetTaskStream(worker, workers.LoadTaskStream, context.Background()) + workerMgr.SetTaskStream(worker, workers.ComputeTaskStream, context.Background()) + oldCtx, cancelOld := context.WithCancel(context.Background()) + newCtx, cancelNew := context.WithCancel(context.Background()) + defer cancelOld() + defer cancelNew() + paused, release, published := make(chan struct{}), make(chan struct{}), make(chan struct{}) + oldStream := &publicationTestStream{ + ctx: metadata.NewIncomingContext(oldCtx, metadata.Pairs("worker_name", worker.Name)), + onRead: func(read int32) { + if read == 2 { + close(paused) + <-release + } + }, + } + newStream := &publicationTestStream{ + ctx: metadata.NewIncomingContext(newCtx, metadata.Pairs("worker_name", worker.Name)), + onRead: func(read int32) { + if read == 3 { + close(published) + } + }, + } + oldDone, newDone := make(chan error, 1), make(chan error, 1) + go func() { oldDone <- new(ServerHandler).SuperStep(&pb.SuperStepReq{}, oldStream) }() + <-paused + started := make(chan struct{}) + go func() { + close(started) + newDone <- new(ServerHandler).SuperStep(&pb.SuperStepReq{}, newStream) + }() + <-started + select { + case <-published: + t.Error("replacement registration overtook an unpublished dispatch entry") + case <-time.After(50 * time.Millisecond): + } + close(release) + <-published + cancelOld() + if err := <-oldDone; err != nil { + t.Error(err) + } + if !workerMgr.IsWorkerReady(worker) || ServerMgr.SuperStepServer(worker.Name).streamServer != newStream { + t.Error("old publication or cleanup replaced the live stream's readiness") + } + cancelNew() + if err := <-newDone; err != nil { + t.Error(err) + } + if workerMgr.IsWorkerReady(worker) { + t.Error("a canceled replacement must no longer be ready") + } +} + +func (s *serializedLoadTestStream) Send(*pb.LoadGraphTaskResp) error { + active := s.active.Add(1) + defer s.active.Add(-1) + for peak := s.peak.Load(); active > peak; peak = s.peak.Load() { + if s.peak.CompareAndSwap(peak, active) { + break + } + } + if s.calls.Add(1) == 1 { + close(s.entered) + <-s.release + return s.sendErr + } + close(s.second) + return nil +} + +func TestWorkerAvailabilitySerializesLoadSends(t *testing.T) { + for _, failed := range []bool{false, true} { + t.Run(fmt.Sprintf("first_send_fails_%v", failed), func(t *testing.T) { + stream := &serializedLoadTestStream{ + entered: make(chan struct{}), release: make(chan struct{}), second: make(chan struct{}), + } + if failed { + stream.sendErr = errors.New("fixture send failure") + } + server := &LoadGraphServer{streamServer: stream} + first, second := make(chan error, 1), make(chan error, 1) + go func() { first <- server.SendLoadReq(&LoadingTaskReq{TaskId: 1, Step: pb.LoadStep_Vertex}) }() + <-stream.entered + started := make(chan struct{}) + go func() { + close(started) + second <- server.AsyncLoad(1, pb.LoadStep_Error, "", "graph", "space", nil, nil) + }() + <-started + select { + case <-stream.second: + t.Error("a second underlying Send entered before the first completed") + case <-time.After(50 * time.Millisecond): + } + close(stream.release) + if err := <-first; !errors.Is(err, stream.sendErr) { + t.Fatalf("original send error changed: %v", err) + } + if err := <-second; err != nil { + t.Fatal(err) + } + if stream.peak.Load() != 1 || stream.calls.Load() != 2 { + t.Fatal("all load sends must be serialized and the lock released after errors") + } + }) + } +} + +func (w *registrationWindowWorker) RuntimeAction(context.Context, *pb.RuntimeActionReq) (*pb.RuntimeActionResp, error) { + close(w.started) + <-w.release + return &pb.RuntimeActionResp{Response: &pb.RuntimeActionResp_MemoryLimitResp{ + MemoryLimitResp: &pb.SetMemoryLimitResp{}, + }}, nil +} + +// This executes real SayHelloMaster while its ordinary memory-limit RPC is pending. +// No worker task stream can exist yet: the worker only opens those after hello returns. +func TestWorkerAvailabilityWaitsForRegistration(t *testing.T) { + previousWorkers := common.PrometheusMetrics.WorkerCnt + previousTasks := common.PrometheusMetrics.TaskCnt + previousRunning := common.PrometheusMetrics.TaskRunningCnt + t.Cleanup(func() { + common.PrometheusMetrics.WorkerCnt = previousWorkers + common.PrometheusMetrics.TaskCnt = previousTasks + common.PrometheusMetrics.TaskRunningCnt = previousRunning + }) + common.PrometheusMetrics.WorkerCnt = prometheus.NewGaugeVec(prometheus.GaugeOpts{Name: "test_registration_workers"}, nil) + common.PrometheusMetrics.TaskCnt = prometheus.NewCounterVec( + prometheus.CounterOpts{Name: "test_registration_tasks"}, []string{"type"}) + common.PrometheusMetrics.TaskRunningCnt = prometheus.NewGaugeVec( + prometheus.GaugeOpts{Name: "test_registration_running"}, []string{"type"}) + workers.WorkerManager.Init() + structure.SpaceManager.Init() + graphMgr.Init() + taskMgr.Init() + if err := taskMgr.InitStore(); err != nil { + t.Fatal(err) + } + ServerMgr.Init() + Scheduler.resourceManager = &schedules.SchedulerResourceManager{} + Scheduler.resourceManager.Init() + Scheduler.taskManager = new(schedules.SchedulerTaskManager).Init() + Scheduler.algorithmManager = &schedules.SchedulerAlgorithmManager{} + Scheduler.algorithmManager.Init() + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + w := ®istrationWindowWorker{started: make(chan struct{}), release: make(chan struct{})} + server := grpc.NewServer() + pb.RegisterWorkerServer(server, w) + go server.Serve(listener) + defer server.Stop() + returned := make(chan error, 1) + ctx := peer.NewContext(context.Background(), &peer.Peer{Addr: listener.Addr()}) + go func() { + _, err := new(ServerHandler).SayHelloMaster(ctx, &pb.HelloMasterReq{ + WorkerPeer: listener.Addr().String(), Version: version.Version, + }) + returned <- err + }() + select { + case <-w.started: + case <-time.After(3 * time.Second): + t.Fatal("registration did not reach memory RPC") + } + defer func() { + close(w.release) + select { + case err := <-returned: + if err != nil { + t.Error(err) + } + case <-time.After(time.Second): + t.Error("hello did not return") + } + }() + if workers.WorkerManager.HasAvailableWorkers("current", nil) { + t.Fatal("hello still pending: reverse dial must not advertise task readiness") + } + task, err := taskMgr.CreateTask("current", structure.TaskTypeLoad, 0) + if err != nil { + t.Fatal(err) + } + task.GraphName = "graph" + task.CreateType = structure.TaskCreateSync + if _, err := taskMgr.AddTask(task); err != nil { + t.Fatal(err) + } + agent, status, _, err := new(schedules.Broker).Init().ApplyAgent(task) + if err != nil || agent != nil || status != schedules.AgentStatusWorkerNotReady { + t.Fatalf("scheduler admitted a worker before its task streams: %v %v", status, err) + } + registered := workers.WorkerManager.GetAllWorkers() + if len(registered) != 1 { + t.Fatal("wrong registration count") + } + worker := registered[0] + t.Cleanup(worker.Close) + t.Run("superstep_publication_order", func(t *testing.T) { testSuperStepPublication(t, worker) }) + if ServerMgr.LoadServer(worker.Name) != nil { + t.Fatal("unexpected load stream") + } + err = sendLoadingReq(task, registered, func(*workers.WorkerClient) *LoadingTaskReq { + return &LoadingTaskReq{} + }) + if err == nil { + t.Fatal("missing task transport must return an error") + } + if task.State != structure.TaskStateError || task.ErrMessage != err.Error() { + t.Fatalf("failed synchronous dispatch left state=%s, message=%q", task.State, task.ErrMessage) + } + persisted, readErr := taskMgr.ReadTask(task.ID) + if readErr != nil { + t.Fatal(readErr) + } + if persisted.State != structure.TaskStateError || persisted.ErrMessage != task.ErrMessage { + t.Fatalf("persisted failed task has state=%s, message=%q", persisted.State, persisted.ErrMessage) + } + + // Reuse the isolated task store to check failures after earlier workers accepted a load. + for _, test := range []struct { + name string + missing bool + failedSend bool + failedStop bool + }{ + {name: "dispatch_missing_stream", missing: true}, + {name: "dispatch_send_failure", failedSend: true}, + {name: "dispatch_cleanup_failure", missing: true, failedStop: true}, + {name: "dispatch_success"}, + } { + t.Run(test.name, func(t *testing.T) { + ServerMgr.Init() + defer ServerMgr.Init() + first, second, third := &dispatchTestStream{}, &dispatchTestStream{}, &dispatchTestStream{} + dispatchErr := errors.New("dispatch failed") + if test.failedSend { + third.sendErr = dispatchErr + } + if test.failedStop { + first.stopErr = errors.New("cleanup failed") + } + ServerMgr.loadSevers["first"] = &LoadGraphServer{streamServer: first} + ServerMgr.loadSevers["second"] = &LoadGraphServer{streamServer: second} + if !test.missing { + ServerMgr.loadSevers["third"] = &LoadGraphServer{streamServer: third} + } + task, err := taskMgr.CreateTask("dispatch-space", structure.TaskTypeLoad, 0) + if err != nil { + t.Fatal(err) + } + task.GraphName = "dispatch-graph" + if _, err := taskMgr.AddTask(task); err != nil { + t.Fatal(err) + } + err = sendLoadingReq(task, []*workers.WorkerClient{ + {Name: "first"}, {Name: "second"}, {Name: "third"}, + }, func(*workers.WorkerClient) *LoadingTaskReq { + return &LoadingTaskReq{TaskId: task.ID, Step: pb.LoadStep_Vertex, + SpaceName: task.SpaceName, GraphName: task.GraphName} + }) + failed := test.missing || test.failedSend + if (err != nil) != failed { + t.Fatalf("dispatch error = %v, want failed = %v", err, failed) + } + if test.failedSend && !errors.Is(err, dispatchErr) { + t.Fatalf("original dispatch error was replaced: %v", err) + } + if failed { + persisted, readErr := taskMgr.ReadTask(task.ID) + if readErr != nil { + t.Fatal(readErr) + } + if task.State != structure.TaskStateError || task.ErrMessage != err.Error() || + persisted.State != structure.TaskStateError || persisted.ErrMessage != err.Error() { + t.Fatal("partial dispatch failure must retain the original error in memory and storage") + } + } + for _, stream := range []*dispatchTestStream{first, second} { + want := 1 + if failed { + want = 2 + } + if len(stream.requests) != want || stream.requests[0].Step != pb.LoadStep_Vertex { + t.Fatalf("successful worker received %d messages, want %d", len(stream.requests), want) + } + if failed { + stop := stream.requests[1] + if stop.Step != pb.LoadStep_Error || stop.TaskId != task.ID || + stop.SpaceName != task.SpaceName || stop.GraphName != task.GraphName { + t.Fatalf("wrong cleanup scope: %v", stop) + } + } + } + }) + } + + t.Run("compute_dispatch", testPartialComputeDispatch) +} diff --git a/vermeer/apps/master/schedules/broker.go b/vermeer/apps/master/schedules/broker.go index 435219dc9..1493d561d 100644 --- a/vermeer/apps/master/schedules/broker.go +++ b/vermeer/apps/master/schedules/broker.go @@ -127,22 +127,15 @@ func (b *Broker) ApplyAgent(taskInfo *structure.TaskInfo, forceApply ...bool) (* // } func (b *Broker) isWorkersReady(workers map[string]*workers.WorkerClient) bool { - ok := false + if len(workers) == 0 { + return false + } for _, w := range workers { - if w.Connection == nil { - ok = false - break - } - - if w.Connection.GetState().String() == "READY" { - ok = true - } else { - ok = false - break + if !workerMgr.IsWorkerReady(w) { + return false } } - - return ok + return true } func (b *Broker) isAgentBusy(agent *Agent) bool { diff --git a/vermeer/apps/master/services/http_worker_availability.go b/vermeer/apps/master/services/http_worker_availability.go new file mode 100644 index 000000000..bdccd7342 --- /dev/null +++ b/vermeer/apps/master/services/http_worker_availability.go @@ -0,0 +1,51 @@ +/* +Licensed to the Apache Software Foundation (ASF) under one or more +contributor license agreements. See the NOTICE file distributed with this +work for additional information regarding copyright ownership. The ASF +licenses this file to You under the Apache License, Version 2.0 (the +"License"); you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +License for the specific language governing permissions and limitations +under the License. +*/ + +package services + +import ( + "net/http" + "vermeer/apps/common" + + "github.com/gin-gonic/gin" +) + +// WorkerAvailabilityHandler reports readiness only for the authenticated space. +type WorkerAvailabilityHandler struct { + common.SenHandler +} + +func (wh *WorkerAvailabilityHandler) GET(ctx *gin.Context) { + writeWorkerAvailability(ctx, WorkerMgr.HasAvailableWorkers) +} + +func writeWorkerAvailability(ctx *gin.Context, available func(string, []string) bool) { + space := credential(ctx).Space() + ctx.JSON(http.StatusOK, gin.H{"available": available(space, availabilityGraphNames(space))}) +} + +func availabilityGraphNames(space string) []string { + // Match AddGraph's publication lock before GetGraphs takes the space lock. + // Names are immutable; do not read the graph's changing Workers slice. + defer GraphMgr.Unlock(GraphMgr.Lock()) + graphs := GraphMgr.GetGraphs(space) + names := make([]string, 0, len(graphs)) + for _, graph := range graphs { + names = append(names, graph.Name) + } + return names +} diff --git a/vermeer/apps/master/services/http_worker_availability_test.go b/vermeer/apps/master/services/http_worker_availability_test.go new file mode 100644 index 000000000..0a5bb9b1d --- /dev/null +++ b/vermeer/apps/master/services/http_worker_availability_test.go @@ -0,0 +1,135 @@ +/* +Licensed to the Apache Software Foundation (ASF) under one or more +contributor license agreements. See the NOTICE file distributed with this +work for additional information regarding copyright ownership. The ASF +licenses this file to You under the Apache License, Version 2.0 (the +"License"); you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +License for the specific language governing permissions and limitations +under the License. +*/ + +package services + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "vermeer/apps/auth" + "vermeer/apps/common" + "vermeer/apps/structure" + + "github.com/gin-gonic/gin" + "github.com/prometheus/client_golang/prometheus" +) + +func TestWorkerAvailabilityUsesAuthenticatedSpace(t *testing.T) { + credential, err := structure.NewCredential("space-user", "own-space") + if err != nil { + t.Fatal(err) + } + previousGraphs := common.PrometheusMetrics.GraphCnt + common.PrometheusMetrics.GraphCnt = prometheus.NewGaugeVec( + prometheus.GaugeOpts{Name: "test_availability_graphs"}, nil) + GraphMgr.Init() + t.Cleanup(func() { GraphMgr.Init(); common.PrometheusMetrics.GraphCnt = previousGraphs }) + for _, space := range []string{"own-space", "another-space"} { + if err := GraphMgr.AddGraph(GraphMgr.CreateGraph(space, space+"-graph")); err != nil { + t.Fatal(err) + } + } + router := gin.New() + var checked string + router.GET("/workers/availability", func(ctx *gin.Context) { + ctx.Set(structure.CredentialKey, credential) + writeWorkerAvailability(ctx, func(space string, graphs []string) bool { + checked = space + return space == "own-space" && len(graphs) == 1 && graphs[0] == "own-space-graph" + }) + }) + response := httptest.NewRecorder() + router.ServeHTTP(response, httptest.NewRequest(http.MethodGet, + "/workers/availability?space=another-space", nil)) + if checked != "own-space" { + t.Fatalf("checked untrusted space %q", checked) + } + if response.Code != http.StatusOK { + t.Fatalf("status = %d", response.Code) + } + var body map[string]any + if err := json.Unmarshal(response.Body.Bytes(), &body); err != nil { + t.Fatal(err) + } + if len(body) != 1 || body["available"] != true { + t.Fatalf("response must contain only readiness: %v", body) + } +} + +func TestWorkerAvailabilityRequiresAuthentication(t *testing.T) { + previousFactor := common.GetConfig("auth_token_factor") + previousMetrics := common.PrometheusMetrics + previousRegisterer, previousGatherer := prometheus.DefaultRegisterer, prometheus.DefaultGatherer + previousMode := gin.Mode() + registry := prometheus.NewRegistry() + prometheus.DefaultRegisterer, prometheus.DefaultGatherer = registry, registry + const factor = "synthetic-worker-availability-factor" + common.SetConfig("auth_token_factor", factor) + auth.Init() + GraphMgr.Init() + t.Cleanup(func() { + common.SetConfig("auth_token_factor", previousFactor) + common.PrometheusMetrics = previousMetrics + prometheus.DefaultRegisterer, prometheus.DefaultGatherer = previousRegisterer, previousGatherer + gin.SetMode(previousMode) + GraphMgr.Init() + }) + factory := auth.TokenFactory{Factor: factor} + token := factory.NewToken("space-user", "own-space", "") + factory.Sign(token) + encoded, err := json.Marshal(token) + if err != nil { + t.Fatal(err) + } + sentinel := &common.Sentinel{} + sentinel.Init("127.0.0.1:0", "release") + SetRouters(sentinel, auth.TokenFilter) + for _, path := range []string{"/workers/availability", "/api/v1/workers/availability"} { + for _, test := range []struct { + name string + header string + status int + }{ + {name: "missing token", status: http.StatusUnauthorized}, + {name: "invalid token", header: "Bearer invalid", status: http.StatusUnauthorized}, + {name: "signed token", header: "Bearer " + auth.ToBase58(string(encoded)), status: http.StatusOK}, + } { + t.Run(path+"/"+test.name, func(t *testing.T) { + req := httptest.NewRequest(http.MethodGet, path, nil) + if test.header != "" { + req.Header.Set(auth.HeaderKey, test.header) + } + response := httptest.NewRecorder() + sentinel.ServeHTTP(response, req) + if response.Code != test.status { + t.Fatalf("status = %d, want %d", response.Code, test.status) + } + if test.status == http.StatusOK { + var body map[string]any + if err := json.Unmarshal(response.Body.Bytes(), &body); err != nil { + t.Fatal(err) + } + if len(body) != 1 || body["available"] != false { + t.Fatalf("unexpected readiness response: %v", body) + } + } + }) + } + } +} diff --git a/vermeer/apps/master/services/router.go b/vermeer/apps/master/services/router.go index ac01e8340..2c3887ec5 100644 --- a/vermeer/apps/master/services/router.go +++ b/vermeer/apps/master/services/router.go @@ -63,6 +63,9 @@ func SetRouters(sen *common.Sentinel, authFilters ...gin.HandlerFunc) { regVerAPI(sen, 1, "/workers", map[string]common.BaseHandler{ "": &WorkersHandler{}, }) + regVerAPI(sen, 1, "/workers", map[string]common.BaseHandler{ + "/availability": &WorkerAvailabilityHandler{}, + }, authFilters...) // /master regVerAPI(sen, 1, "/master", map[string]common.BaseHandler{ diff --git a/vermeer/apps/master/workers/worker_availability_test.go b/vermeer/apps/master/workers/worker_availability_test.go new file mode 100644 index 000000000..c4b49ca5a --- /dev/null +++ b/vermeer/apps/master/workers/worker_availability_test.go @@ -0,0 +1,207 @@ +/* +Licensed to the Apache Software Foundation (ASF) under one or more +contributor license agreements. See the NOTICE file distributed with this +work for additional information regarding copyright ownership. The ASF +licenses this file to You under the Apache License, Version 2.0 (the +"License"); you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +License for the specific language governing permissions and limitations +under the License. +*/ + +package workers + +import ( + "context" + "maps" + "net" + "sync" + "testing" + "time" + + "google.golang.org/grpc" + "google.golang.org/grpc/connectivity" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/test/bufconn" +) + +func readyWorkerConnection(t *testing.T) *grpc.ClientConn { + t.Helper() + listener := bufconn.Listen(1024 * 1024) + server := grpc.NewServer() + go func() { _ = server.Serve(listener) }() + t.Cleanup(func() { server.Stop(); _ = listener.Close() }) + connection, err := grpc.NewClient("passthrough:///availability-test", + grpc.WithTransportCredentials(insecure.NewCredentials()), + grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { return listener.Dial() })) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = connection.Close() }) + connection.Connect() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + for { + state := connection.GetState() + if state == connectivity.Ready { + return connection + } + if !connection.WaitForStateChange(ctx, state) { + t.Fatal("worker connection did not become ready") + } + } +} + +func readyTaskStreams() [taskStreamCount]context.Context { + return [taskStreamCount]context.Context{ + context.Background(), context.Background(), context.Background(), + } +} + +func TestWorkerAvailabilityRespectsSpaceAllocation(t *testing.T) { + connection := readyWorkerConnection(t) + cases := []struct { + name string + group string + mappings map[string]string + graphs []string + available bool + }{ + {"unallocated named group", "default", nil, nil, false}, + {"another space", "default", map[string]string{"other": "default"}, nil, false}, + {"space allocation", "default", map[string]string{"current": "default"}, nil, true}, + {"graph allocation", "default", map[string]string{"current graph": "default"}, []string{"graph"}, true}, + {"other graph", "default", map[string]string{"current-other graph": "default"}, []string{"graph"}, false}, + {"common worker", "$", nil, nil, true}, + {"other space with shared prefix", "default", map[string]string{"current graph": "default"}, nil, false}, + } + for _, test := range cases { + t.Run(test.name, func(t *testing.T) { + manager := workerManager{ + workersByName: map[string]*WorkerClient{ + "worker": {Name: "worker", Group: test.group, Connection: connection, + taskStreams: readyTaskStreams()}, + }, + groupMapper: test.mappings, + } + before := maps.Clone(manager.groupMapper) + if got := manager.HasAvailableWorkers("current", test.graphs); got != test.available { + t.Fatalf("availability = %v, want %v", got, test.available) + } + if !maps.Equal(manager.groupMapper, before) { + t.Fatal("readiness changed allocations") + } + }) + } +} + +func TestWorkerAvailabilityHandlesRegistrationAndDisconnect(t *testing.T) { + connection := readyWorkerConnection(t) + manager := workerManager{workersByName: map[string]*WorkerClient{}, groupMapper: map[string]string{}} + if manager.HasAvailableWorkers("current", nil) { + t.Fatal("missing worker must not be ready") + } + manager.workersByName["worker"] = &WorkerClient{Name: "worker", Group: "$"} + if manager.HasAvailableWorkers("current", nil) { + t.Fatal("unconnected worker must not be ready") + } + worker := manager.workersByName["worker"] + manager.SetWorkerConnection(worker, connection) + if manager.HasAvailableWorkers("current", nil) { + t.Fatal("reverse dial alone must not advertise a task-ready worker") + } + for stream := LoadTaskStream; stream < taskStreamCount; stream++ { + manager.SetTaskStream(worker, stream, context.Background()) + if stream < SuperStepStream && manager.HasAvailableWorkers("current", nil) { + t.Fatal("partially registered task streams must not be ready") + } + } + if !manager.HasAvailableWorkers("current", nil) { + t.Fatal("registered ready worker must be available") + } + var wait sync.WaitGroup + wait.Add(1) + go func() { + defer wait.Done() + for i := 0; i < 1000; i++ { + manager.HasAvailableWorkers("current", nil) + } + }() + for i := 0; i < 1000; i++ { + manager.SetWorkerConnection(worker, nil) + manager.SetWorkerConnection(worker, connection) + } + _ = connection.Close() + wait.Wait() + if manager.HasAvailableWorkers("current", nil) { + t.Fatal("disconnected worker must not be ready") + } +} + +func TestWorkerAvailabilityWaitsForSelectedGroup(t *testing.T) { + connection := readyWorkerConnection(t) + manager := workerManager{ + workersByName: map[string]*WorkerClient{ + "first": {Name: "first", Group: "assigned", Connection: connection, taskStreams: readyTaskStreams()}, + "second": {Name: "second", Group: "assigned", taskStreams: readyTaskStreams()}, + "common": {Name: "common", Group: "$", Connection: connection, taskStreams: readyTaskStreams()}, + }, + groupMapper: map[string]string{"current": "assigned"}, + } + if manager.HasAvailableWorkers("current", nil) { + t.Fatal("common workers cannot override an explicit group with an unready member") + } + manager.SetWorkerConnection(manager.workersByName["second"], connection) + if !manager.HasAvailableWorkers("current", nil) { + t.Fatal("the selected group must become available when all members are ready") + } + delete(manager.workersByName, "first") + delete(manager.workersByName, "second") + if manager.HasAvailableWorkers("current", nil) { + t.Fatal("an empty explicit group cannot fall back to common workers") + } + manager.workersByName["graph-worker"] = &WorkerClient{ + Name: "graph-worker", Group: "graph-group", Connection: connection, taskStreams: readyTaskStreams(), + } + manager.groupMapper[manager.toOwnerKey("current", "graph")] = "graph-group" + if !manager.HasAvailableWorkers("current", []string{"graph"}) { + t.Fatal("an existing graph's explicit ready group remains an available workload") + } +} + +func TestWorkerAvailabilityTracksTaskStreamCancellationAndReplacement(t *testing.T) { + worker := &WorkerClient{Name: "worker", Group: "$", Connection: readyWorkerConnection(t), + taskStreams: readyTaskStreams()} + manager := workerManager{workersByName: map[string]*WorkerClient{"worker": worker}} + previous, cancel := context.WithCancel(context.Background()) + manager.SetTaskStream(worker, LoadTaskStream, previous) + if !manager.IsWorkerReady(worker) { + t.Fatal("complete worker must be schedulable") + } + cancel() + if manager.IsWorkerReady(worker) || manager.HasAvailableWorkers("current", nil) { + t.Fatal("a closed task stream must disable scheduling and advertised availability") + } + replacement, closeReplacement := context.WithCancel(context.Background()) + defer closeReplacement() + manager.SetTaskStream(worker, LoadTaskStream, replacement) + manager.ClearTaskStream(worker, LoadTaskStream, previous) + if !manager.IsWorkerReady(worker) { + t.Fatal("old stream cleanup must not clear its replacement") + } + manager.ClearTaskStream(worker, LoadTaskStream, replacement) + if manager.IsWorkerReady(worker) { + t.Fatal("cleared task stream must not be schedulable") + } + manager.workersByName[worker.Name] = &WorkerClient{Name: worker.Name} + manager.SetTaskStream(worker, LoadTaskStream, replacement) + if manager.IsWorkerReady(worker) { + t.Fatal("a removed worker instance must not be schedulable") + } +} diff --git a/vermeer/apps/master/workers/worker_manager.go b/vermeer/apps/master/workers/worker_manager.go index f2ea2fb62..53773c3fc 100644 --- a/vermeer/apps/master/workers/worker_manager.go +++ b/vermeer/apps/master/workers/worker_manager.go @@ -18,6 +18,7 @@ under the License. package workers import ( + "context" "encoding/json" "fmt" "math/rand" @@ -40,6 +41,15 @@ const ( prefixMapper = "WORKER_MAPPER" ) +type TaskStream uint8 + +const ( + LoadTaskStream TaskStream = iota + ComputeTaskStream + SuperStepStream + taskStreamCount +) + // WorkerClient type WorkerClient struct { Id int32 `json:"id,omitempty"` @@ -54,8 +64,9 @@ type WorkerClient struct { //LoadStream pb.Master_LoadGraphTaskServer `json:"-"` //ComputeStream pb.Master_ComputeTaskServer `json:"-"` //SuperStepStream pb.Master_SuperStepServer `json:"-"` - Connection *grpc.ClientConn `json:"-"` - Session pb.WorkerClient `json:"-"` + Connection *grpc.ClientConn `json:"-"` + Session pb.WorkerClient `json:"-"` + taskStreams [taskStreamCount]context.Context } func (wc *WorkerClient) Close() { @@ -233,6 +244,79 @@ func (wm *workerManager) CheckWorkerAlive(name string) bool { return true } +// SetWorkerConnection publishes a completed registration dial to readiness readers. +func (wm *workerManager) SetWorkerConnection(worker *WorkerClient, connection *grpc.ClientConn) { + defer wm.Unlock(wm.Lock()) + worker.Connection = connection +} + +// Task streams are published only after their dispatch entry has been installed. +func (wm *workerManager) SetTaskStream(worker *WorkerClient, stream TaskStream, ctx context.Context) { + defer wm.Unlock(wm.Lock()) + if worker != nil && wm.workersByName[worker.Name] == worker { + worker.taskStreams[stream] = ctx + } +} + +func (wm *workerManager) ClearTaskStream(worker *WorkerClient, stream TaskStream, ctx context.Context) { + defer wm.Unlock(wm.Lock()) + if worker != nil && wm.workersByName[worker.Name] == worker && worker.taskStreams[stream] == ctx { + worker.taskStreams[stream] = nil + } +} + +func (wm *workerManager) IsWorkerReady(worker *WorkerClient) bool { + defer wm.Unlock(wm.Lock()) + return wm.workerReadyLocked(worker) +} + +func (wm *workerManager) workerReadyLocked(worker *WorkerClient) bool { + if worker == nil || wm.workersByName[worker.Name] != worker || + worker.Connection == nil || worker.Connection.GetState() != connectivity.Ready { + return false + } + for _, ctx := range worker.taskStreams { + if ctx == nil || ctx.Err() != nil { + return false + } + } + return true +} + +// HasAvailableWorkers follows existing allocation precedence without allocating. +// Every registered member needs both a ready reverse connection and live task streams. +func (wm *workerManager) HasAvailableWorkers(space string, graphNames []string) bool { + defer wm.Unlock(wm.Lock()) + groups := make(map[string]bool) + group, explicit := wm.groupMapper[space] + if !explicit { + group = "$" // AddWorker normalizes ungrouped registered workers to "$". + } + groups[group] = true + for _, graph := range graphNames { + if group, ok := wm.groupMapper[wm.toOwnerKey(space, graph)]; ok { + groups[group] = true + } + } + for group := range groups { + found, ready := false, true + for _, worker := range wm.workersByName { + if worker.Group != group { + continue + } + found = true + if !wm.workerReadyLocked(worker) { + ready = false + break + } + } + if found && ready { + return true + } + } + return false +} + func (wm *workerManager) SetWorkerGroup(workerName, workerGroup string) error { if workerName == "" { return fmt.Errorf("the argument 'workerName' is invalid") diff --git a/vermeer/apps/structure/graph.go b/vermeer/apps/structure/graph.go index 142ede959..19a19866b 100644 --- a/vermeer/apps/structure/graph.go +++ b/vermeer/apps/structure/graph.go @@ -89,11 +89,11 @@ func (vg *VermeerGraph) MallocData(option GraphDataBackendOption) { vg.BackendOption = option vg.Data.MallocData(option) } -func (vg *VermeerGraph) SetOption(useOutEdges bool, useOutDegree bool, useProperty bool) { +func (vg *VermeerGraph) SetOption(useOutEdges bool, useOutDegree bool, useProperty bool) error { vg.UseOutEdges = useOutEdges vg.UseOutDegree = useOutDegree vg.UseProperty = useProperty - vg.Data.SetOption(GraphDataOption{ + return vg.Data.SetOption(GraphDataOption{ graphName: vg.Name, spaceName: vg.SpaceName, dataDir: vg.DataDir, @@ -340,7 +340,7 @@ func (vg *VermeerGraph) Read(workerName string) error { } data := new(GraphData) data.MallocData(vg.BackendOption) - data.SetOption(GraphDataOption{ + err = data.SetOption(GraphDataOption{ graphName: vg.Name, spaceName: vg.SpaceName, dataDir: vg.DataDir, @@ -348,6 +348,9 @@ func (vg *VermeerGraph) Read(workerName string) error { useOutEdges: vg.UseOutEdges, useOutDegree: vg.UseOutDegree, }) + if err != nil { + return err + } err = data.Load(vg.DataDir) if err != nil { return err diff --git a/vermeer/apps/structure/graph_data.go b/vermeer/apps/structure/graph_data.go index 92a5acbc7..80b5e9e17 100644 --- a/vermeer/apps/structure/graph_data.go +++ b/vermeer/apps/structure/graph_data.go @@ -25,7 +25,7 @@ import ( ) type VertexInterface interface { - Init(dataDir string) + Init(dataDir string) error TotalVertexCount() uint32 GetVertex(vertexID uint32) Vertex GetVertexIndex(vertex string) (uint32, bool) @@ -144,13 +144,16 @@ func (gd *GraphData) MallocData(option GraphDataBackendOption) { gd.InEdgesProperty = &EdgeProperties{} } -func (gd *GraphData) SetOption(option GraphDataOption) { +func (gd *GraphData) SetOption(option GraphDataOption) error { gd.graphName = option.graphName gd.spaceName = option.spaceName if option.firstInit { - gd.Vertex.Init(option.dataDir) + if err := gd.Vertex.Init(option.dataDir); err != nil { + return err + } } gd.Edges.Init(option.useOutEdges, option.useOutDegree) + return nil } //func (gd *GraphData) Name() string { diff --git a/vermeer/apps/structure/property.go b/vermeer/apps/structure/property.go index 94f6f23d0..e4a9aa31f 100644 --- a/vermeer/apps/structure/property.go +++ b/vermeer/apps/structure/property.go @@ -483,6 +483,10 @@ func GetSchemaFromHugegraph(params map[string]string) (PropertySchema, PropertyS hgName := options.GetString(params, "load.hugegraph_name") username := options.GetString(params, "load.hugegraph_username") password := options.GetString(params, "load.hugegraph_password") + username, password, err := common.ResolveHugegraphCredentials(username, password) + if err != nil { + return PropertySchema{}, PropertySchema{}, err + } useProperty := options.GetInt(params, "load.use_property") == 1 vertexPropIDs := strings.Split(options.GetString(params, "load.vertex_property"), ",") edgePropIDs := strings.Split(options.GetString(params, "load.edge_property"), ",") diff --git a/vermeer/apps/structure/property_credentials_test.go b/vermeer/apps/structure/property_credentials_test.go new file mode 100644 index 000000000..38a36ab4a --- /dev/null +++ b/vermeer/apps/structure/property_credentials_test.go @@ -0,0 +1,139 @@ +/* +Licensed to the Apache Software Foundation (ASF) under one or more +contributor license agreements. See the NOTICE file distributed with this +work for additional information regarding copyright ownership. The ASF +licenses this file to You under the Apache License, Version 2.0 (the +"License"); you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package structure + +import ( + "context" + "net" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "reflect" + "strings" + "sync/atomic" + "testing" + + "google.golang.org/grpc" + "vermeer/apps/common" + "vermeer/apps/options" + pd "vermeer/apps/protos/hugegraph-pd-grpc" +) + +type schemaCredentialPD struct { + pd.UnimplementedPDServer + pd.UnimplementedDiscoveryServiceServer + serverURL string +} + +func (s *schemaCredentialPD) GetMembers(context.Context, *pd.GetMembersRequest) (*pd.GetMembersResponse, error) { + return &pd.GetMembersResponse{}, nil +} + +func (s *schemaCredentialPD) GetNodes(context.Context, *pd.Query) (*pd.NodeInfos, error) { + return &pd.NodeInfos{Info: []*pd.NodeInfo{{ + Address: s.serverURL, Labels: map[string]string{"GRAPHSPACE": "DEFAULT"}, + }}}, nil +} + +func TestHugegraphSchemaDeploymentCredentials(t *testing.T) { + previousOptions := options.AllOptionsMap + options.Init() + t.Cleanup(func() { options.AllOptionsMap = previousOptions }) + var authorized atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + user, password, ok := r.BasicAuth() + if !ok || user != "schema-reader" || password != "deployment-secret" { + w.WriteHeader(http.StatusUnauthorized) + return + } + if r.URL.Path != "/graphspaces/DEFAULT/graphs/fresh-graph/schema" || r.URL.Query().Get("format") != "json" { + w.WriteHeader(http.StatusNotFound) + return + } + authorized.Add(1) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"vertexlabels":[],"edgelabels":[],"propertykeys":[]}`)) + })) + defer server.Close() + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + discovery := grpc.NewServer() + fixture := &schemaCredentialPD{serverURL: server.URL} + pd.RegisterPDServer(discovery, fixture) + pd.RegisterDiscoveryServiceServer(discovery, fixture) + go func() { _ = discovery.Serve(listener) }() + defer discovery.Stop() + credentials := filepath.Join(t.TempDir(), "credentials.json") + if err := os.WriteFile(credentials, []byte(`{"username":"schema-reader","password":"deployment-secret"}`), 0600); err != nil { + t.Fatal(err) + } + t.Setenv("VERMEER_HUGEGRAPH_CREDENTIALS_FILE", credentials) + t.Setenv("VERMEER_PD_PEERS", `["`+listener.Addr().String()+`"]`) + params := map[string]string{ + "load.hg_pd_peers": `["` + listener.Addr().String() + `"]`, + "load.hugegraph_name": "DEFAULT/fresh-graph/g", "load.use_property": "1", + "load.hugegraph_username": "wrong-task-user", "load.hugegraph_password": "wrong-task-password", + } + original := make(map[string]string) + for k, v := range params { + original[k] = v + } + vertex, edge, err := GetSchemaFromHugegraph(params) + if err != nil { + t.Fatal(err) + } + if vertex.HgPSchema == nil || edge.HgPSchema == nil { + t.Fatal("authenticated Schema response must initialize vertex and edge schemas") + } + if authorized.Load() != 2 { + t.Fatalf("expected authenticated discovery probe and Schema read, got %d", authorized.Load()) + } + if !reflect.DeepEqual(params, original) { + t.Fatal("deployment credentials must not be written into task params") + } +} + +func TestHugegraphSchemaRejectsInvalidCredentialsBeforeDiscovery(t *testing.T) { + t.Setenv("VERMEER_HUGEGRAPH_CREDENTIALS_FILE", filepath.Join(t.TempDir(), "missing.json")) + // There is no PD endpoint: credential validation must fail before discovery. + if _, _, err := GetSchemaFromHugegraph(map[string]string{}); err == nil || !strings.Contains(err.Error(), "read HugeGraph credentials file") { + t.Fatalf("expected credential error before discovery, got %v", err) + } +} + +func TestHugegraphSchemaRequiresDeploymentPD(t *testing.T) { + previousOptions := options.AllOptionsMap + options.Init() + t.Cleanup(func() { options.AllOptionsMap = previousOptions }) + previousPeers := common.GetConfigDefault("pd_peers", "") + common.SetConfig("pd_peers", "") + t.Cleanup(func() { common.SetConfig("pd_peers", previousPeers) }) + t.Setenv("VERMEER_PD_PEERS", "") + credentials := filepath.Join(t.TempDir(), "credentials.json") + if err := os.WriteFile(credentials, []byte(`{"username":"fixture-reader","password":"fixture-password"}`), 0600); err != nil { + t.Fatal(err) + } + t.Setenv("VERMEER_HUGEGRAPH_CREDENTIALS_FILE", credentials) + _, _, err := GetSchemaFromHugegraph(map[string]string{}) + if err == nil || !strings.Contains(err.Error(), "configure pd_peers") { + t.Fatal("deployment Schema credentials require deployment PD configuration before discovery") + } +} diff --git a/vermeer/apps/structure/vertex.go b/vermeer/apps/structure/vertex.go index 33dbef181..411a65ffc 100644 --- a/vermeer/apps/structure/vertex.go +++ b/vermeer/apps/structure/vertex.go @@ -183,9 +183,10 @@ type VertexMem struct { VertexLongIDMap map[string]uint32 } -func (vm *VertexMem) Init(dataDir string) { +func (vm *VertexMem) Init(dataDir string) error { vm.TotalVertex = make([]Vertex, 0) _ = dataDir + return nil } func (vm *VertexMem) TotalVertexCount() uint32 { @@ -252,11 +253,9 @@ type VertexInDB struct { dataDir string } -func (vm *VertexInDB) Init(dataDir string) { - var err error +func (vm *VertexInDB) Init(dataDir string) error { vm.dataDir = vm.makeDataDir(dataDir) - vm.sliceStores = make([]storage.Store, 1) - vm.sliceStores[0], err = storage.StoreMaker(storage.StoreOption{ + store, err := storage.StoreMaker(storage.StoreOption{ StoreName: storage.StoreTypePebble, Path: path.Join(vm.dataDir, "self_vertex_slice"), Fsync: false, @@ -264,9 +263,10 @@ func (vm *VertexInDB) Init(dataDir string) { UseFilter: false, }) if err != nil { - logrus.Errorf("init vertex store error:%v", err) - return + return fmt.Errorf("initialize vertex store: %w", err) } + vm.sliceStores = []storage.Store{store} + return nil } func (vm *VertexInDB) makeDataDir(dataDir string) string { diff --git a/vermeer/apps/structure/vertex_init_test.go b/vermeer/apps/structure/vertex_init_test.go new file mode 100644 index 000000000..8b60a3327 --- /dev/null +++ b/vermeer/apps/structure/vertex_init_test.go @@ -0,0 +1,74 @@ +/* +Licensed to the Apache Software Foundation (ASF) under one or more +contributor license agreements. See the NOTICE file distributed with this +work for additional information regarding copyright ownership. The ASF +licenses this file to You under the Apache License, Version 2.0 (the +"License"); you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +License for the specific language governing permissions and limitations +under the License. +*/ + +package structure + +import ( + "os" + "path/filepath" + "strings" + "testing" +) + +func blockedVertexDirectory(t *testing.T) string { + t.Helper() + file := filepath.Join(t.TempDir(), "not-a-directory") + if err := os.WriteFile(file, []byte("fixture"), 0600); err != nil { + t.Fatal(err) + } + return file +} + +func TestVertexStoreInitializationErrorAndRetry(t *testing.T) { + vertex := &VertexInDB{} + err := vertex.Init(blockedVertexDirectory(t)) + if err == nil || !strings.Contains(err.Error(), "initialize vertex store") { + t.Fatalf("expected vertex storage initialization error, got %v", err) + } + // Failed initialization must remain safe to unload before a retry. + vertex.freeMem() + if err := vertex.Init(t.TempDir()); err != nil { + t.Fatal(err) + } + defer vertex.freeMem() + vertex.workers = []*GraphWorker{{IsSelf: true}} + vertex.AppendVertices(Vertex{ID: "retry-vertex"}) + if vertex.TotalVertexCount() != 1 || vertex.GetVertex(0).ID != "retry-vertex" { + t.Fatal("successful retry must store the loaded vertex") + } +} + +func TestGraphPropagatesVertexStoreInitializationError(t *testing.T) { + graph := &VermeerGraph{DataDir: blockedVertexDirectory(t)} + graph.MallocData(GraphDataBackendOption{VertexDataBackend: DataBackendInDB}) + if err := graph.SetOption(true, true, false); err == nil { + t.Fatal("graph initialization must report vertex storage failure") + } + graph.FreeMem() +} + +func TestMemoryVertexInitialization(t *testing.T) { + graph := &VermeerGraph{} + graph.MallocData(GraphDataBackendOption{VertexDataBackend: DataBackendInMem}) + if err := graph.SetOption(true, true, false); err != nil { + t.Fatal(err) + } + graph.Data.Vertex.AppendVertices(Vertex{ID: "memory-vertex"}) + if graph.Data.Vertex.GetVertex(0).ID != "memory-vertex" { + t.Fatal("memory vertex backend must remain usable") + } +} diff --git a/vermeer/apps/worker/compute_bl.go b/vermeer/apps/worker/compute_bl.go index 2e5a5cb50..bbfa74e82 100644 --- a/vermeer/apps/worker/compute_bl.go +++ b/vermeer/apps/worker/compute_bl.go @@ -117,7 +117,26 @@ func (cb *ComputeBl) StartCompute(taskID int32, spaceName string, graphName stri return } - ct.Task.SetState(structure.TaskStateInitOK) + for { + ct.Locker.Lock() + action := atomic.LoadInt32(&ct.Task.Action) + if action != structure.ActionPauseTask { + if action == structure.ActionCancelTask { + cb.cancelAction(ct) + ct.Locker.Unlock() + return + } + ct.Task.SetState(structure.TaskStateInitOK) + ct.WaitingForStep = true + ct.Locker.Unlock() + break + } + // Pause outside the handoff lock, then recheck before publishing readiness. + ct.Locker.Unlock() + if !cb.pauseAction(ct) { + return + } + } req := pb.ComputeTaskStatusReq{ TaskId: taskID, State: string(structure.TaskStateInitOK), @@ -127,8 +146,13 @@ func (cb *ComputeBl) StartCompute(taskID int32, spaceName string, graphName stri } _, err = ServiceWorker.MasterClient.ComputeTaskStatus(context.Background(), &req) if err != nil { - logrus.Errorf("ComputeTaskStatus error:%v", err) - cb.SetStatusError(taskID, fmt.Sprintf("ComputeTaskStatus error:%v", err)) + ct.Locker.Lock() + ct.WaitingForStep = false + ct.Locker.Unlock() + if cb.CheckAction(ct) { + logrus.Errorf("ComputeTaskStatus error:%v", err) + cb.SetStatusError(taskID, fmt.Sprintf("ComputeTaskStatus error:%v", err)) + } } //go cb.RunSuperStep(ct.Task.ID, nil) @@ -144,6 +168,12 @@ func (cb *ComputeBl) RunSuperStep(taskID int32, computeValues map[string][]byte) } }() computeTask := ComputeTaskMgr.GetTask(taskID) + if computeTask == nil { + return + } + computeTask.Locker.Lock() + computeTask.WaitingForStep = false + computeTask.Locker.Unlock() if !cb.CheckAction(computeTask) { return } @@ -764,9 +794,11 @@ func (cb *ComputeBl) CheckAction(computeTask *compute.ComputerTask) (isContinue } func (cb *ComputeBl) cancelAction(computeTask *compute.ComputerTask) { - computeTask.Task.SetState(structure.TaskStateCanceled) - common.PrometheusMetrics.TaskRunningCnt.WithLabelValues(computeTask.Task.Type).Dec() - time.AfterFunc(1*time.Minute, func() { cb.endTask(computeTask.Task.ID) }) + computeTask.CancelOnce.Do(func() { + computeTask.Task.SetState(structure.TaskStateCanceled) + common.PrometheusMetrics.TaskRunningCnt.WithLabelValues(computeTask.Task.Type).Dec() + time.AfterFunc(1*time.Minute, func() { cb.endTask(computeTask.Task.ID) }) + }) } func (cb *ComputeBl) pauseAction(computeTask *compute.ComputerTask) bool { diff --git a/vermeer/apps/worker/compute_dispatch_cancellation_test.go b/vermeer/apps/worker/compute_dispatch_cancellation_test.go new file mode 100644 index 000000000..0eabf0cd6 --- /dev/null +++ b/vermeer/apps/worker/compute_dispatch_cancellation_test.go @@ -0,0 +1,317 @@ +/* +Licensed to the Apache Software Foundation (ASF) under one or more +contributor license agreements. See the NOTICE file distributed with this +work for additional information regarding copyright ownership. The ASF +licenses this file to You under the Apache License, Version 2.0 (the +"License"); you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package worker + +import ( + "bytes" + "context" + "errors" + "runtime" + "testing" + "time" + "vermeer/apps/common" + "vermeer/apps/compute" + "vermeer/apps/options" + pb "vermeer/apps/protos" + "vermeer/apps/structure" + + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/testutil" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +type computeStatusClient struct { + pb.MasterClient + calls int + entered chan struct{} + release chan struct{} + reportErr error +} + +func (c *computeStatusClient) ComputeTaskStatus(context.Context, *pb.ComputeTaskStatusReq, + ...grpc.CallOption) (*pb.ComputeTaskStatusResp, error) { + c.calls++ + if c.entered != nil && c.calls == 1 { + close(c.entered) + <-c.release + } + return &pb.ComputeTaskStatusResp{}, c.reportErr +} + +type cancelInitWorker struct { + compute.WorkerComputerBase + entered chan struct{} + release chan struct{} +} + +func (w *cancelInitWorker) Init() error { + close(w.entered) + <-w.release + return nil +} + +type cancelInitMaker struct { + compute.AlgorithmMakerBase + worker *cancelInitWorker +} + +func (m *cancelInitMaker) CreateWorkerComputer() compute.WorkerComputer { return m.worker } + +func TestComputeDispatchCancellation(t *testing.T) { + previousOptions := options.AllOptionsMap + previousGraphs := common.PrometheusMetrics.GraphCnt + previousTasks := common.PrometheusMetrics.TaskCnt + previousRunning := common.PrometheusMetrics.TaskRunningCnt + previousMaster := ServiceWorker.MasterClient + options.Init() + options.AllOptionsMap["compute.algorithm"] = &options.Option{ + Name: "compute.algorithm", ValueType: options.OptionValueTypeString, Value: "base", Convert: options.ToString, + } + common.PrometheusMetrics.GraphCnt = prometheus.NewGaugeVec(prometheus.GaugeOpts{Name: "test_compute_cancel_graphs"}, nil) + common.PrometheusMetrics.TaskCnt = prometheus.NewCounterVec(prometheus.CounterOpts{Name: "test_compute_cancel_tasks"}, []string{"type"}) + common.PrometheusMetrics.TaskRunningCnt = prometheus.NewGaugeVec(prometheus.GaugeOpts{Name: "test_compute_cancel_running"}, []string{"type"}) + GraphMgr.Init() + TaskMgr.Init() + ComputeTaskMgr.Init() + AlgorithmMgr.Init() + t.Cleanup(func() { + GraphMgr.Init() + TaskMgr.Init() + ComputeTaskMgr.Init() + options.AllOptionsMap = previousOptions + common.PrometheusMetrics.GraphCnt = previousGraphs + common.PrometheusMetrics.TaskCnt = previousTasks + common.PrometheusMetrics.TaskRunningCnt = previousRunning + ServiceWorker.MasterClient = previousMaster + }) + graph := GraphMgr.CreateGraph("compute-cancel", "graph") + graph.Workers = []*structure.GraphWorker{{Name: "worker", IsSelf: true}} + if err := GraphMgr.AddGraph(graph); err != nil { + t.Fatal(err) + } + handler := new(PeerHandler) + t.Run("before_task_creation", func(t *testing.T) { + _, err := handler.ControlTask(context.Background(), &pb.ControlTaskReq{ + TaskID: 7001, Action: structure.ActionCancelTask}) + if status.Code(err) != codes.NotFound { + t.Fatalf("missing task cancellation must be retryable NotFound, got %v", err) + } + }) + for index := 0; index < 3; index++ { + duringInit := index == 0 + name := "waiting_for_first_step" + if duringInit { + name = "during_initialization" + } else if index == 2 { + name = "status_error_after_cancel" + } + t.Run(name, func(t *testing.T) { + taskID := int32(7002 + index) + statusClient := new(computeStatusClient) + if index == 2 { + statusClient.entered = make(chan struct{}) + statusClient.release = make(chan struct{}) + statusClient.reportErr = errors.New("status report failed") + } + ServiceWorker.MasterClient = statusClient + initWorker := &cancelInitWorker{entered: make(chan struct{}), release: make(chan struct{})} + if duringInit { + AlgorithmMgr.Register(&cancelInitMaker{worker: initWorker}, "test") + } else { + AlgorithmMgr.Register(new(compute.AlgorithmMakerBase), "test") + } + done := make(chan struct{}) + go func() { + defer close(done) + new(ComputeBl).StartCompute(taskID, graph.SpaceName, graph.Name, "base", map[string]string{"compute.parallel": "1"}) + }() + if duringInit { + select { + case <-initWorker.entered: + case <-time.After(time.Second): + t.Fatal("worker did not reach initialization") + } + } else if index == 2 { + select { + case <-statusClient.entered: + case <-time.After(time.Second): + t.Fatal("status report was not reached") + } + } else { + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("initialization did not finish") + } + } + _, err := handler.ControlTask(context.Background(), &pb.ControlTaskReq{ + TaskID: taskID, Action: structure.ActionCancelTask}) + if err != nil { + t.Fatal(err) + } + if duringInit { + close(initWorker.release) + } else if index == 2 { + close(statusClient.release) + } + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("canceled initialization did not finish") + } + ct := ComputeTaskMgr.GetTask(taskID) + if ct == nil || ct.Task.State != structure.TaskStateCanceled { + t.Fatalf("cancellation was not consumed: %+v", ct) + } + wantReports := 1 + if duringInit { + wantReports = 0 + } + if statusClient.calls != wantReports { + t.Errorf("status reports=%d want=%d", statusClient.calls, wantReports) + } + // An already queued first-step message must not consume cancellation twice. + new(ComputeBl).RunSuperStep(taskID, nil) + if ct.Task.State != structure.TaskStateCanceled { + t.Fatal("late first step resumed a canceled task") + } + if running := testutil.ToFloat64(common.PrometheusMetrics.TaskRunningCnt.WithLabelValues(structure.TaskTypeCompute)); running != 0 { + t.Fatalf("cancellation left running gauge=%v", running) + } + }) + } + for index, phase := range []string{"initialization", "first_step", "status_error"} { + t.Run("paused_"+phase, func(t *testing.T) { + taskID := int32(7010 + index) + statusClient := new(computeStatusClient) + if phase == "status_error" { + statusClient.entered = make(chan struct{}) + statusClient.release = make(chan struct{}) + statusClient.reportErr = errors.New("status report failed") + } + ServiceWorker.MasterClient = statusClient + initWorker := &cancelInitWorker{entered: make(chan struct{}), release: make(chan struct{})} + if phase == "initialization" { + AlgorithmMgr.Register(&cancelInitMaker{worker: initWorker}, "test") + } else { + AlgorithmMgr.Register(new(compute.AlgorithmMakerBase), "test") + } + start := func() { + new(ComputeBl).StartCompute(taskID, graph.SpaceName, graph.Name, "base", + map[string]string{"compute.parallel": "1"}) + } + done := make(chan struct{}) + var entered <-chan struct{} + if phase == "initialization" { + entered = initWorker.entered + } else if phase == "status_error" { + entered = statusClient.entered + } + if entered == nil { + start() + } else { + go func() { defer close(done); start() }() + select { + case <-entered: + case <-time.After(time.Second): + t.Fatal("compute did not reach the controlled initialization gate") + } + } + if _, err := handler.ControlTask(context.Background(), &pb.ControlTaskReq{ + TaskID: taskID, Action: structure.ActionPauseTask}); err != nil { + t.Fatal(err) + } + switch phase { + case "initialization": + close(initWorker.release) + case "status_error": + close(statusClient.release) + default: + go func() { defer close(done); new(ComputeBl).RunSuperStep(taskID, nil) }() + } + // Observe the actual pause wait before canceling, without changing its 10-second cadence. + if !waitForComputePause() { + t.Error("compute did not enter the pause wait") + } + canceled := make(chan error, 1) + go func() { + _, err := handler.ControlTask(context.Background(), &pb.ControlTaskReq{ + TaskID: taskID, Action: structure.ActionCancelTask}) + canceled <- err + }() + select { + case err := <-canceled: + if err != nil { + t.Error(err) + } + case <-time.After(time.Second): + t.Error("cancel RPC blocked behind the paused computation") + // Join the old-source negative control too, before restoring shared fixtures. + select { + case <-canceled: + case <-time.After(11 * time.Second): + t.Fatal("cancel RPC did not finish after the pause poll") + } + } + select { + case <-done: + case <-time.After(11 * time.Second): + t.Fatal("paused computation did not consume cancellation") + } + ct := ComputeTaskMgr.GetTask(taskID) + if ct.Task.State != structure.TaskStateCanceled { + t.Fatalf("paused computation ended in state %s", ct.Task.State) + } + wantReports := 1 + if phase == "initialization" { + wantReports = 0 + } + if statusClient.calls != wantReports { + t.Errorf("status reports=%d want=%d", statusClient.calls, wantReports) + } + new(ComputeBl).RunSuperStep(taskID, nil) + if ct.Task.State != structure.TaskStateCanceled { + t.Error("late first step resumed a canceled paused task") + } + if running := testutil.ToFloat64(common.PrometheusMetrics.TaskRunningCnt.WithLabelValues(structure.TaskTypeCompute)); running != 0 { + t.Errorf("paused cancellation left running gauge=%v", running) + } + }) + } +} + +func waitForComputePause() bool { + stack := make([]byte, 128<<10) + deadline := time.NewTimer(time.Second) + defer deadline.Stop() + tick := time.NewTicker(time.Millisecond) + defer tick.Stop() + for { + n := runtime.Stack(stack, true) + if bytes.Contains(stack[:n], []byte("vermeer/apps/worker.(*ComputeBl).pauseAction(")) { + return true + } + select { + case <-deadline.C: + return false + case <-tick.C: + } + } +} diff --git a/vermeer/apps/worker/grpc_handlers.go b/vermeer/apps/worker/grpc_handlers.go index 6cd4fe17a..87d38b05d 100644 --- a/vermeer/apps/worker/grpc_handlers.go +++ b/vermeer/apps/worker/grpc_handlers.go @@ -26,10 +26,13 @@ import ( "vermeer/apps/common" "vermeer/apps/graphio" pb "vermeer/apps/protos" + "vermeer/apps/structure" "github.com/sirupsen/logrus" + "google.golang.org/grpc/codes" "google.golang.org/grpc/metadata" "google.golang.org/grpc/peer" + "google.golang.org/grpc/status" ) type PeerHandler struct { @@ -155,12 +158,22 @@ func (ph *PeerHandler) ControlTask(ctx context.Context, req *pb.ControlTaskReq) _ = ctx task := TaskMgr.GetTaskByID(req.GetTaskID()) if task == nil { - return &pb.ControlTaskResp{}, nil + return nil, status.Error(codes.NotFound, "task is not initialized") } err := TaskMgr.SetAction(task, req.Action) if err != nil { return &pb.ControlTaskResp{}, err } + if req.Action == structure.ActionCancelTask && task.Type == structure.TaskTypeCompute { + if computeTask := ComputeTaskMgr.GetTask(task.ID); computeTask != nil { + computeTask.Locker.Lock() + if computeTask.WaitingForStep { + computeTask.WaitingForStep = false + new(ComputeBl).CheckAction(computeTask) + } + computeTask.Locker.Unlock() + } + } //SELECT: // for { // switch task.State { diff --git a/vermeer/apps/worker/load_graph_bl.go b/vermeer/apps/worker/load_graph_bl.go index 36aae49af..9f0dfade4 100644 --- a/vermeer/apps/worker/load_graph_bl.go +++ b/vermeer/apps/worker/load_graph_bl.go @@ -39,13 +39,24 @@ type LoadGraphBl struct { } func (lb *LoadGraphBl) StartLoadGraph(spaceName string, taskId int32, graphName string, workers []string, params map[string]string) { + var initializing *graphio.LoadGraphTask + var initializationWg *sync.WaitGroup + var initializationError string + initialized := false defer func() { if r := recover(); r != nil { - lb.SetStatusError(taskId, fmt.Sprintf("StartLoadGraph panic recover panic:%v, stack message: %s", r, - common.GetCurrentGoroutineStack())) - logrus.Errorf("StartLoadGraph panic recover taskID:%v, panic:%v, stack message: %s", taskId, r, + initializationError = fmt.Sprintf("StartLoadGraph panic recover panic:%v, stack message: %s", r, common.GetCurrentGoroutineStack()) } + // Release local receivers before status reporting, which may block on Master. + if initializing != nil { + initializing.InitializationFailed = !initialized + initializationWg.Done() + } + if initializationError != "" { + logrus.Errorf("initialize graph task %d failed: %s", taskId, initializationError) + lb.SetStatusError(taskId, initializationError) + } }() graph := GraphMgr.CreateGraph(spaceName, graphName) @@ -80,8 +91,9 @@ func (lb *LoadGraphBl) StartLoadGraph(spaceName string, taskId int32, graphName TaskMgr.AddTask(task) logrus.Infof("create load task: %d, graph: %s", taskId, graphName) loadTask := LoadGraphMgr.InstallTask(task) - // recv vertex wait until init done - loadTask.LoadWg.Add(1) + // InstallTask holds the barrier before publishing the task to receivers. + initializing = loadTask + initializationWg = loadTask.LoadWg ctx := context.Background() parallel := options.GetInt(params, "load.parallel") if parallel <= 0 { @@ -93,8 +105,8 @@ func (lb *LoadGraphBl) StartLoadGraph(spaceName string, taskId int32, graphName _, err = graph.SetDataDir(spaceName, graphName, ServiceWorker.WorkerName) if err != nil { - logrus.Errorf("set data dir error: %s", err) - lb.SetStatusError(taskId, fmt.Sprintf("set data dir error: %s", err)) + initializationError = fmt.Sprintf("set data dir error: %s", err) + return } // remove history data graph.Remove() @@ -120,8 +132,7 @@ func (lb *LoadGraphBl) StartLoadGraph(spaceName string, taskId int32, graphName //logrus.Infof(" hugegraph vertex schema %v", graph.Data.VertexPropertySchema) //logrus.Infof(" hugegraph edge schema %v", graph.Data.InEdgesPropertySchema) if err != nil { - lb.SetStatusError(taskId, fmt.Sprintf("load schema from hugegraph error:%v", err)) - logrus.Errorf("load schema from hugegraph error:%v", err) + initializationError = fmt.Sprintf("load schema from hugegraph error:%v", err) return } @@ -130,12 +141,15 @@ func (lb *LoadGraphBl) StartLoadGraph(spaceName string, taskId int32, graphName graph.Data.VertexPropertySchema.Init(options.GetMapString(params, "load.vertex_property")) graph.Data.InEdgesPropertySchema.Init(options.GetMapString(params, "load.edge_property")) } - graph.SetOption(useOutEdges, useOutDegree, useProperty) + if err := graph.SetOption(useOutEdges, useOutDegree, useProperty); err != nil { + initializationError = fmt.Sprintf("initialize graph data: %s", err) + return + } if graph.UseProperty { graph.Data.VertexProperty.Init(graph.Data.VertexPropertySchema) } - loadTask.LoadWg.Done() + initialized = true for i := 0; i < parallel; i++ { valueCtx := context.WithValue(ctx, "worker_id", i) go lb.RunLoadVertex(valueCtx, loadTask) @@ -385,6 +399,9 @@ func (lb *LoadGraphBl) GatherVertex(taskID int32, workerName string, end bool, e //defer lb.CheckAction(loadTask) graph := GraphMgr.GetGraphByName(loadTask.Task.SpaceName, loadTask.Task.GraphName) loadTask.LoadWg.Wait() + if loadTask.InitializationFailed { + return + } loadTask.RecvWg.Add(1) //gw := graph.GetGraphWorker(workerName) @@ -867,6 +884,10 @@ func (lb *LoadGraphBl) RecvVertex(taskId int32, worker string, count int32, end } loadTask.RecvWg.Add(1) loadTask.LoadWg.Wait() + if loadTask.InitializationFailed { + loadTask.RecvWg.Done() + return + } //defer lb.CheckAction(loadTask) graph := GraphMgr.GetGraphByName(loadTask.Task.SpaceName, loadTask.Task.GraphName) atomic.AddUint32(&graph.Data.VertexCount, uint32(count)) @@ -967,6 +988,9 @@ func (lb *LoadGraphBl) RecvEdge(taskId int32, worker string, count int32, end bo //defer lb.CheckAction(loadTask) graph := GraphMgr.GetGraphByName(loadTask.Task.SpaceName, loadTask.Task.GraphName) loadTask.LoadWg.Wait() + if loadTask.InitializationFailed { + return + } loadTask.RecvWg.Add(1) logrus.Debugf("recv edge count: %d, end: %v,endNum:%v, worker: %s", count, end, endNum, worker) diff --git a/vermeer/apps/worker/load_graph_readiness_test.go b/vermeer/apps/worker/load_graph_readiness_test.go new file mode 100644 index 000000000..36e7a68ce --- /dev/null +++ b/vermeer/apps/worker/load_graph_readiness_test.go @@ -0,0 +1,240 @@ +/* +Licensed to the Apache Software Foundation (ASF) under one or more +contributor license agreements. See the NOTICE file distributed with this +work for additional information regarding copyright ownership. The ASF +licenses this file to You under the Apache License, Version 2.0 (the +"License"); you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package worker + +import ( + "context" + "sync" + "sync/atomic" + "testing" + "time" + + "vermeer/apps/common" + "vermeer/apps/graphio" + "vermeer/apps/options" + pb "vermeer/apps/protos" + "vermeer/apps/structure" + + "github.com/prometheus/client_golang/prometheus" + "google.golang.org/grpc" +) + +type blockingLoadStatusClient struct { + pb.MasterClient + started chan struct{} + release <-chan struct{} +} + +func (client *blockingLoadStatusClient) LoadTaskStatus(context.Context, *pb.LoadTaskStatusReq, + ...grpc.CallOption) (*pb.LoadTaskStatusResp, error) { + close(client.started) + <-client.release + return &pb.LoadTaskStatusResp{}, nil +} + +func receiveFixture(t *testing.T) *graphio.LoadGraphTask { + t.Helper() + GraphMgr.Init() + LoadGraphMgr.Init() + task := LoadGraphMgr.InstallTask(&structure.TaskInfo{ + ID: 5001, SpaceName: "readiness-test", GraphName: "graph", + Workers: []*structure.TaskWorker{{Name: "worker"}}, + }) + t.Cleanup(func() { GraphMgr.Init(); LoadGraphMgr.DeleteTask(task.Task.ID) }) + return task +} + +func TestLoadReceiversUseInitializationResult(t *testing.T) { + task := receiveFixture(t) + task.InitializationFailed = true + task.LoadWg.Done() + // Runtime task errors can change after initialization. Receivers must only + // inspect the immutable result published by the initialization barrier. + stopped, started := make(chan struct{}), make(chan struct{}) + var stop atomic.Bool + go func() { + defer close(stopped) + task.Task.State = structure.TaskStateError + close(started) + for !stop.Load() { + task.Task.State = structure.TaskStateError + } + }() + <-started + t.Cleanup(func() { stop.Store(true); <-stopped }) + receiver := &LoadGraphBl{} + for i := 0; i < 100; i++ { + receiver.RecvVertex(task.Task.ID, "worker", 0, false, 0, nil) + receiver.GatherVertex(task.Task.ID, "worker", false, 0, nil) + receiver.RecvEdge(task.Task.ID, "worker", 0, false, 0, nil) + } + if count := atomic.LoadInt32((*int32)(task.RecvWg)); count != 0 { + t.Fatalf("failed initialization leaked %d registered receivers", count) + } +} + +func TestLoadReceiversBarrierSurvivesTaskCleanup(t *testing.T) { + GraphMgr.Init() + LoadGraphMgr.Init() + TaskMgr.Init() + previousOptions := options.AllOptionsMap + options.Init() + previousGraphs := common.PrometheusMetrics.GraphCnt + common.PrometheusMetrics.GraphCnt = prometheus.NewGaugeVec( + prometheus.GaugeOpts{Name: "test_cleanup_graphs"}, nil) + t.Cleanup(func() { + GraphMgr.Init() + LoadGraphMgr.DeleteTask(-5002) + TaskMgr.Init() + options.AllOptionsMap = previousOptions + common.PrometheusMetrics.GraphCnt = previousGraphs + }) + var task *graphio.LoadGraphTask + var barrier *sync.WaitGroup + options.AllOptionsMap["load.type"] = &options.Option{ + Convert: func(string) (interface{}, error) { + task = LoadGraphMgr.GetLoadTask(-5002) + barrier = task.LoadWg + // Simulate cancellation cleanup while initialization is still active. + task.FreeMemory() + panic("initialization failed after cleanup") + }, + } + reported, release, returned := make(chan struct{}), make(chan struct{}), make(chan struct{}) + previousMaster := ServiceWorker.MasterClient + ServiceWorker.MasterClient = &blockingLoadStatusClient{started: reported, release: release} + t.Cleanup(func() { + close(release) + select { + case <-returned: + case <-time.After(time.Second): + t.Error("initialization did not return after status report resumed") + } + ServiceWorker.MasterClient = previousMaster + }) + go func() { + defer close(returned) + (&LoadGraphBl{}).StartLoadGraph("readiness-test", -5002, "graph", nil, + map[string]string{"load.type": "cleanup-test"}) + }() + select { + case <-reported: + case <-time.After(time.Second): + t.Fatal("initialization did not report its failure") + } + select { + case <-returned: + t.Fatal("status report must still be blocked") + default: + } + if task == nil || barrier == nil || task.LoadWg != nil || !task.InitializationFailed { + t.Fatal("initialization must publish failure after task cleanup") + } + released := make(chan struct{}) + go func() { barrier.Wait(); close(released) }() + select { + case <-released: + case <-time.After(time.Second): + t.Fatal("blocked status report left the original initialization barrier blocked") + } +} + +func TestRecvVertexRegistersBeforeInitialization(t *testing.T) { + for _, failed := range []bool{false, true} { + name := "success" + if failed { + name = "failure" + } + t.Run(name, func(t *testing.T) { + task := receiveFixture(t) + previous := common.PrometheusMetrics.GraphCnt + common.PrometheusMetrics.GraphCnt = prometheus.NewGaugeVec( + prometheus.GaugeOpts{Name: "test_receive_graphs"}, nil) + t.Cleanup(func() { common.PrometheusMetrics.GraphCnt = previous }) + graph := GraphMgr.CreateGraph(task.Task.SpaceName, task.Task.GraphName) + graph.MallocData(structure.GraphDataBackendOption{VertexDataBackend: structure.DataBackendInMem}) + if err := graph.SetOption(false, false, false); err != nil { + t.Fatal(err) + } + if err := GraphMgr.AddGraph(graph); err != nil { + t.Fatal(err) + } + var once sync.Once + release := func() { + once.Do(func() { + task.InitializationFailed = failed + task.LoadWg.Done() + }) + } + received := make(chan struct{}) + vertex := structure.Vertex{ID: "delayed-vertex"} + data := make([]byte, vertex.PredictSize()) + if _, err := vertex.Marshal(data); err != nil { + t.Fatal(err) + } + go func() { + (&LoadGraphBl{}).RecvVertex(task.Task.ID, "worker", 1, false, 0, data) + close(received) + }() + t.Cleanup(func() { + release() + select { + case <-received: + case <-time.After(time.Second): + t.Error("receiver did not leave initialization wait") + } + }) + deadline := time.Now().Add(time.Second) + for atomic.LoadInt32((*int32)(task.RecvWg)) == 0 && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if count := atomic.LoadInt32((*int32)(task.RecvWg)); count != 1 { + t.Fatalf("receiver must register before waiting for initialization, got %d", count) + } + completion := make(chan struct{}) + go func() { task.RecvWg.Wait(); close(completion) }() + // This is the same completion wait used by an ending vertex frame. + select { + case <-completion: + t.Fatal("completion advanced past an initialization-blocked receiver") + case <-time.After(20 * time.Millisecond): + } + release() + select { + case <-received: + case <-time.After(time.Second): + t.Fatal("receiver did not complete") + } + select { + case <-completion: + case <-time.After(time.Second): + t.Fatal("completion waiter leaked after receiver exit") + } + if atomic.LoadInt32((*int32)(task.RecvWg)) != 0 { + t.Fatal("unbalanced receiver accounting") + } + if failed { + if graph.Data.VertexCount != 0 { + t.Fatal("failed initialization must not accept vertices") + } + } else if graph.Data.VertexCount != 1 || graph.Data.Vertex.GetVertex(0).ID != vertex.ID { + t.Fatal("successful initialization must process the registered vertex") + } + }) + } +} diff --git a/vermeer/docs/hubble-integration/2026-10-23/01-vermeer-loaded-pagerank.png b/vermeer/docs/hubble-integration/2026-10-23/01-vermeer-loaded-pagerank.png new file mode 100644 index 000000000..53c480ea6 Binary files /dev/null and b/vermeer/docs/hubble-integration/2026-10-23/01-vermeer-loaded-pagerank.png differ diff --git a/vermeer/docs/hubble-integration/2026-10-23/02-vermeer-tasks-success.png b/vermeer/docs/hubble-integration/2026-10-23/02-vermeer-tasks-success.png new file mode 100644 index 000000000..95a1fe3d5 Binary files /dev/null and b/vermeer/docs/hubble-integration/2026-10-23/02-vermeer-tasks-success.png differ diff --git a/vermeer/docs/hubble-integration/2026-10-23/03-vermeer-pagerank-results.png b/vermeer/docs/hubble-integration/2026-10-23/03-vermeer-pagerank-results.png new file mode 100644 index 000000000..45e599547 Binary files /dev/null and b/vermeer/docs/hubble-integration/2026-10-23/03-vermeer-pagerank-results.png differ diff --git a/vermeer/docs/hubble-integration/2026-10-23/README.md b/vermeer/docs/hubble-integration/2026-10-23/README.md new file mode 100644 index 000000000..be63a07a8 --- /dev/null +++ b/vermeer/docs/hubble-integration/2026-10-23/README.md @@ -0,0 +1,43 @@ + + +# Hubble integration evidence — 2026-10-23 + +Real Chromium screenshots from a local retained development stack: + +- Hubble: latest remote master `eef5afd29762f394667494b726afdce1b938214d`. +- Vermeer executable: `273e1ecf63bc739068ea08e4e2b9ff47334d94a5`. +- Synthetic graph: `DEFAULT/vermeer_review_20261023`, six vertices and seven directed edges. +- Existing anonymous development mode (`auth=none`); this run does not verify token authorization. + +The Hubble helper's API smoke test completed load task `-13` and 20-step PageRank task `-14`. +Chromium submitted reload task `-15` and PageRank task `-16` with the UI defaults +(10 steps, damping 0.85). Both completed successfully; the task list and full result page rendered +all six finite ranks without page script errors. Independent PageRank iteration on the fixture +matched both outputs within `3e-8` absolute error. + +1. [Loaded graph and PageRank parameters](01-vermeer-loaded-pagerank.png). + Batch results appear in Async Tasks; the interactive canvas is unused in this flow. +2. [Successful load and compute tasks](02-vermeer-tasks-success.png). +3. [Full PageRank result](03-vermeer-pagerank-results.png). + +Before validation, the retained Store exhausted its 32GB disk through 1GB Raft segments despite +its advertised 64MB option. Unused zero extents were reclaimed while Store was stopped, with all +35 segment SHA-256 checksums unchanged. A supported JRaft JVM property limited new segments to +16MB for this local environment; an actual new 16MB segment was observed. This is a local runtime +workaround. No Store binary or Toolchain source change is included here, and this run does not +claim the Store configuration propagation bug has been fixed upstream.