Skip to content

Commit

Permalink
fix(milvus): racy container setup (#2693)
Browse files Browse the repository at this point in the history 
Fix race in container setup caused by copying the files into the
container manually.

While here update tests to use require so they are easier to follow.
  • Loading branch information
@stevenh
stevenh authored Aug 7, 2024
1 parent 8f90dcf commit f333902
Show file tree
Hide file tree
Showing 4 changed files with 39 additions and 60 deletions.
4 changes: 4 additions & 0 deletions modules/milvus/go.mod
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ go 1.21

require (
github.com/milvus-io/milvus-sdk-go/v2 v2.3.6
github.com/stretchr/testify v1.9.0
github.com/testcontainers/testcontainers-go v0.32.0
)

Expand All @@ -19,6 +20,7 @@ require (
github.com/containerd/log v0.1.0 // indirect
github.com/containerd/platforms v0.2.1 // indirect
github.com/cpuguy83/dockercfg v0.3.1 // indirect
github.com/davecgh/go-spew v1.1.1 // indirect
github.com/distribution/reference v0.6.0 // indirect
github.com/docker/docker v27.1.0+incompatible // indirect
github.com/docker/go-connections v0.5.0 // indirect
Expand Down Expand Up @@ -47,6 +49,7 @@ require (
github.com/opencontainers/go-digest v1.0.0 // indirect
github.com/opencontainers/image-spec v1.1.0 // indirect
github.com/pkg/errors v0.9.1 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c // indirect
github.com/rogpeppe/go-internal v1.8.1 // indirect
github.com/shirou/gopsutil/v3 v3.23.12 // indirect
Expand All @@ -69,6 +72,7 @@ require (
google.golang.org/genproto/googleapis/rpc v0.0.0-20240318140521-94a12d6c2237 // indirect
google.golang.org/grpc v1.64.1 // indirect
google.golang.org/protobuf v1.33.0 // indirect
gopkg.in/yaml.v3 v3.0.1 // indirect
)

replace github.com/testcontainers/testcontainers-go => ../..
5 changes: 4 additions & 1 deletion modules/milvus/go.sum
View file
Original file line number Diff line number Diff line change
Expand Up @@ -270,8 +270,9 @@ github.com/spf13/viper v1.3.2/go.mod h1:ZiWeW+zYFKm7srdB9IoDzzZXaJaI5eL9QjNiN/DM
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
github.com/stretchr/objx v0.5.0 h1:1zr/of2m5FGMsad5YfcqgdqdWrIhu+EBEJRhR1U7z/c=
github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
Expand Down Expand Up @@ -482,6 +483,8 @@ google.golang.org/protobuf v1.33.0/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHh
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI=
gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys=
gopkg.in/go-playground/assert.v1 v1.2.1/go.mod h1:9RXL0bg/zibRAgZUYszZSwO/z8Y/a8bDuhia5mkpMnE=
Expand Down
67 changes: 23 additions & 44 deletions modules/milvus/milvus.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ import (
_ "embed"
"fmt"
"html/template"
"os"
"io"
"time"

"github.com/testcontainers/testcontainers-go"
Expand All @@ -16,7 +16,11 @@ import (
//go:embed mounts/embedEtcd.yaml.tpl
var embedEtcdConfigTpl string

const embedEtcdContainerPath string = "/milvus/configs/embedEtcd.yaml"
const (
embedEtcdContainerPath = "/milvus/configs/embedEtcd.yaml"
defaultClientPort = 2379
grpcPort = "19530/tcp"
)

// MilvusContainer represents the Milvus container type used in the module
type MilvusContainer struct {
Expand All @@ -30,7 +34,7 @@ func (c *MilvusContainer) ConnectionString(ctx context.Context) (string, error)
if err != nil {
return "", err
}
port, err := c.MappedPort(ctx, "19530/tcp")
port, err := c.MappedPort(ctx, grpcPort)
if err != nil {
return "", err
}
Expand All @@ -45,6 +49,11 @@ func RunContainer(ctx context.Context, opts ...testcontainers.ContainerCustomize

// Run creates an instance of the Milvus container type
func Run(ctx context.Context, img string, opts ...testcontainers.ContainerCustomizer) (*MilvusContainer, error) {
config, err := renderEmbedEtcdConfig(defaultClientPort)
if err != nil {
return nil, fmt.Errorf("render config: %w", err)
}

req := testcontainers.ContainerRequest{
Image: img,
ExposedPorts: []string{"19530/tcp", "9091/tcp", "2379/tcp"},
Expand All @@ -54,16 +63,13 @@ func Run(ctx context.Context, img string, opts ...testcontainers.ContainerCustom
"ETCD_CONFIG_PATH": embedEtcdContainerPath,
"COMMON_STORAGETYPE": "local",
},
Cmd: []string{"milvus", "run", "standalone"},
WaitingFor: wait.ForHTTP("/healthz").WithPort("9091").WithStartupTimeout(60 * time.Second).WithPollInterval(30 * time.Second),
LifecycleHooks: []testcontainers.ContainerLifecycleHooks{
{
PostCreates: []testcontainers.ContainerHook{
// Copy the default embed etcd config to container after it's created.
// Otherwise the milvus container will panic on startup.
createDefaultEmbedEtcdConfig,
},
},
Cmd: []string{"milvus", "run", "standalone"},
WaitingFor: wait.ForHTTP("/healthz").
WithPort("9091").
WithStartupTimeout(time.Minute).
WithPollInterval(time.Second),
Files: []testcontainers.ContainerFile{
{ContainerFilePath: embedEtcdContainerPath, Reader: config},
},
}

Expand All @@ -90,7 +96,9 @@ type embedEtcdConfigTplParams struct {
Port int
}

func renderEmbedEtcdConfig(port int) ([]byte, error) {
// renderEmbedEtcdConfig renders the embed etcd config template with the given port
// and returns it as an io.Reader.
func renderEmbedEtcdConfig(port int) (io.Reader, error) {
tplParams := embedEtcdConfigTplParams{
Port: port,
}
Expand All @@ -105,34 +113,5 @@ func renderEmbedEtcdConfig(port int) ([]byte, error) {
return nil, fmt.Errorf("failed to render embed etcd config template: %w", err)
}

return embedEtcdYaml.Bytes(), nil
}

// createDefaultEmbedEtcdConfig creates a default embed etcd config file,
// using the default port 2379 as the advertised port. The file is then copied to the container.
func createDefaultEmbedEtcdConfig(ctx context.Context, c testcontainers.Container) error {
// Otherwise the milvus container will panic on startup.
defaultEmbedEtcdConfig, err := renderEmbedEtcdConfig(2379)
if err != nil {
return fmt.Errorf("failed to render default config: %w", err)
}

tmpDir := os.TempDir()
defaultEmbedEtcdConfigPath := fmt.Sprintf("%s/embedEtcd.yaml", tmpDir)

if err := os.WriteFile(defaultEmbedEtcdConfigPath, defaultEmbedEtcdConfig, 0o644); err != nil {
return fmt.Errorf("failed to write default embed etcd config to a temporary dir: %w", err)
}

if err != nil {
return fmt.Errorf("can't create default embed etcd config: %w", err)
}
defer os.Remove(defaultEmbedEtcdConfigPath)

err = c.CopyFileToContainer(ctx, defaultEmbedEtcdConfigPath, embedEtcdContainerPath, 0o644)
if err != nil {
return fmt.Errorf("can't copy %s to container: %w", defaultEmbedEtcdConfigPath, err)
}

return nil
return &embedEtcdYaml, nil
}
23 changes: 8 additions & 15 deletions modules/milvus/milvus_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"testing"

"github.com/milvus-io/milvus-sdk-go/v2/client"
"github.com/stretchr/testify/require"

"github.com/testcontainers/testcontainers-go/modules/milvus"
)
Expand All @@ -13,35 +14,27 @@ func TestMilvus(t *testing.T) {
ctx := context.Background()

container, err := milvus.Run(ctx, "milvusdb/milvus:v2.3.9")
if err != nil {
t.Fatal(err)
}
require.NoError(t, err)

// Clean up the container after the test is complete
t.Cleanup(func() {
if err := container.Terminate(ctx); err != nil {
t.Fatalf("failed to terminate container: %s", err)
}
err = container.Terminate(ctx)
require.NoError(t, err)
})

t.Run("Connect to Milvus with gRPC", func(tt *testing.T) {
// connectionString {
connectionStr, err := container.ConnectionString(ctx)
// }
if err != nil {
tt.Fatal(err)
}
require.NoError(t, err)

milvusClient, err := client.NewGrpcClient(context.Background(), connectionStr)
if err != nil {
tt.Fatal("failed to connect to Milvus:", err.Error())
}
require.NoError(t, err)

defer milvusClient.Close()

v, err := milvusClient.GetVersion(ctx)
if err != nil {
tt.Fatal("failed to get version:", err.Error())
}
require.NoError(t, err)

tt.Logf("Milvus version: %s", v)
})
Expand Down

0 comments on commit f333902

Please sign in to comment.