talos

1
// This Source Code Form is subject to the terms of the Mozilla Public
2
// License, v. 2.0. If a copy of the MPL was not distributed with this
3
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
4

5
package backend
6

7
import (
8
	"crypto/tls"
9

10
	"github.com/siderolabs/gen/concurrent"
11
	"github.com/siderolabs/grpc-proxy/proxy"
12
)
13

14
// APIDFactory caches connection to apid instances by target.
15
//
16
// TODO: need to clean up idle connections from time to time.
17
type APIDFactory struct {
18
	cache    *concurrent.HashTrieMap[string, *APID]
19
	provider TLSConfigProvider
20
}
21

22
// TLSConfigProvider provides tls.Config for client connections.
23
type TLSConfigProvider interface {
24
	ClientConfig() (*tls.Config, error)
25
}
26

27
// NewAPIDFactory creates new APIDFactory with given tls.Config.
28
//
29
// Client TLS config is used to connect to other apid instances.
30
func NewAPIDFactory(provider TLSConfigProvider) *APIDFactory {
31
	return &APIDFactory{
32
		cache:    concurrent.NewHashTrieMap[string, *APID](),
33
		provider: provider,
34
	}
35
}
36

37
// Get backend by target.
38
//
39
// Get performs caching of backends.
40
func (factory *APIDFactory) Get(target string) (proxy.Backend, error) {
41
	b, ok := factory.cache.Load(target)
42
	if ok {
43
		return b, nil
44
	}
45

46
	backend, err := NewAPID(target, factory.provider.ClientConfig)
47
	if err != nil {
48
		return nil, err
49
	}
50

51
	existing, loaded := factory.cache.LoadOrStore(target, backend)
52
	if loaded {
53
		// race: another Get() call built different backend
54
		backend.Close()
55

56
		return existing, nil
57
	}
58

59
	return backend, nil
60
}
61

62
// Flush all cached backends.
63
//
64
// This ensures that all connections are closed.
65
func (factory *APIDFactory) Flush() {
66
	factory.cache.Enumerate(func(key string, backend *APID) bool {
67
		backend.Close()
68

69
		return true
70
	})
71
}
72