weaviate

1
//                           _       _
2
// __      _____  __ ___   ___  __ _| |_ ___
3
// \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
4
//  \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
5
//   \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
6
//
7
//  Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
8
//
9
//  CONTACT: hello@weaviate.io
10
//
11

12
package clients
13

14
import (
15
	"context"
16
	"net/http"
17
	"strings"
18
	"sync"
19
	"time"
20

21
	enterrors "github.com/weaviate/weaviate/entities/errors"
22

23
	"github.com/pkg/errors"
24
	"github.com/weaviate/weaviate/modules/text2vec-transformers/ent"
25
)
26

27
func (v *vectorizer) WaitForStartup(initCtx context.Context,
28
	interval time.Duration,
29
) error {
30
	endpoints := map[string]string{}
31
	if v.originPassage != v.originQuery {
32
		endpoints["passage"] = v.urlPassage("/.well-known/ready", ent.VectorizationConfig{})
33
		endpoints["query"] = v.urlQuery("/.well-known/ready", ent.VectorizationConfig{})
34
	} else {
35
		endpoints[""] = v.urlPassage("/.well-known/ready", ent.VectorizationConfig{})
36
	}
37

38
	ch := make(chan error, len(endpoints))
39
	var wg sync.WaitGroup
40
	for serviceName, endpoint := range endpoints {
41
		serviceName, endpoint := serviceName, endpoint
42
		wg.Add(1)
43
		enterrors.GoWrapper(func() {
44
			defer wg.Done()
45
			if err := v.waitFor(initCtx, interval, endpoint, serviceName); err != nil {
46
				ch <- err
47
			}
48
		}, v.logger)
49
	}
50
	wg.Wait()
51
	close(ch)
52

53
	if len(ch) > 0 {
54
		var errs []string
55
		for err := range ch {
56
			errs = append(errs, err.Error())
57
		}
58
		return errors.Errorf(strings.Join(errs, ", "))
59
	}
60
	return nil
61
}
62

63
func (v *vectorizer) waitFor(initCtx context.Context, interval time.Duration, endpoint string, serviceName string) error {
64
	ticker := time.NewTicker(interval)
65
	defer ticker.Stop()
66
	expired := initCtx.Done()
67
	var lastErr error
68
	prefix := ""
69
	if serviceName != "" {
70
		prefix = "[" + serviceName + "] "
71
	}
72

73
	for {
74
		select {
75
		case <-ticker.C:
76
			lastErr = v.checkReady(initCtx, endpoint, serviceName)
77
			if lastErr == nil {
78
				return nil
79
			}
80
			v.logger.
81
				WithField("action", "transformer_remote_wait_for_startup").
82
				WithError(lastErr).Warnf("%stransformer remote inference service not ready", prefix)
83
		case <-expired:
84
			return errors.Wrapf(lastErr, "%sinit context expired before remote was ready", prefix)
85
		}
86
	}
87
}
88

89
func (v *vectorizer) checkReady(initCtx context.Context, endpoint string, serviceName string) error {
90
	// spawn a new context (derived on the overall context) which is used to
91
	// consider an individual request timed out
92
	// due to parent timeout being superior over request's one, request can be cancelled by parent timeout
93
	// resulting in "send check ready request" even if service is responding with non 2xx http code
94
	requestCtx, cancel := context.WithTimeout(initCtx, 500*time.Millisecond)
95
	defer cancel()
96

97
	req, err := http.NewRequestWithContext(requestCtx, http.MethodGet, endpoint, nil)
98
	if err != nil {
99
		return errors.Wrap(err, "create check ready request")
100
	}
101

102
	res, err := v.httpClient.Do(req)
103
	if err != nil {
104
		return errors.Wrap(err, "send check ready request")
105
	}
106

107
	defer res.Body.Close()
108
	if res.StatusCode > 299 {
109
		return errors.Errorf("not ready: status %d", res.StatusCode)
110
	}
111

112
	return nil
113
}
114