More precise waiting for socket

Signed-off-by: Maksym Pavlenko <[email protected]>

Author: mxpv

go.mod

@@ -6,6 +6,7 @@ require (
 	github.com/go-openapi/strfmt v0.17.1
 	github.com/go-openapi/swag v0.17.1
 	github.com/go-openapi/validate v0.17.1
+	github.com/golang/mock v1.1.1
 	github.com/jessevdk/go-flags v1.4.0
 	github.com/sirupsen/logrus v1.1.1
 	golang.org/x/net v0.0.0-20181023162649-9b4f9f5ad519

go.sum

@@ -35,6 +35,8 @@ github.com/go-openapi/swag v0.17.1/go.mod h1:AByQ+nYG6gQg71GINrmuDXCPWdL640yX49/
 github.com/go-openapi/validate v0.17.0/go.mod h1:Uh4HdOzKt19xGIGm1qHf/ofbX1YQ4Y+MYsct2VUrAJ4=
 github.com/go-openapi/validate v0.17.1 h1:RfQTLHm/gEu0oSUmbTOy0PMufjkE5/pPfnqYpor3WLc=
 github.com/go-openapi/validate v0.17.1/go.mod h1:Uh4HdOzKt19xGIGm1qHf/ofbX1YQ4Y+MYsct2VUrAJ4=
+github.com/golang/mock v1.1.1 h1:G5FRp8JnTd7RQH5kemVNlMeyXQAztQ3mOWV95KxsXH8=
+github.com/golang/mock v1.1.1/go.mod h1:oTYuIxOrZwtPieC+H1uAHpcLFnEyAGVDL/k47Jfbm0A=
 github.com/google/uuid v1.0.0 h1:b4Gk+7WdP/d3HZH8EJsZpvV7EtDOgaZLtnaNGIu1adA=
 github.com/google/uuid v1.0.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/jessevdk/go-flags v1.4.0 h1:4IU2WS7AumrZ/40jfhf4QVDMsQwqA7VEHozFRrGARJA=

machine.go

@@ -299,7 +299,7 @@ func (m *Machine) startVMM(ctx context.Context) (chan error, error) {
 	}()
 
 	// Wait for firecracker to initialize:
-	err = waitForSocket(m.cfg.SocketPath, 3*time.Second, exitchannel)
+	err = m.waitForSocket(3*time.Second, exitchannel)
 	if err != nil {
 		msg := fmt.Sprintf("Firecracker did not create API socket %s: %s", m.cfg.SocketPath, err)
 		err = errors.New(msg)
@@ -514,28 +514,38 @@ func (m *Machine) refreshMachineConfig() error {
 	return nil
 }
 
-// waitFile waits for the given file to exist
-func waitForSocket(path string, timeout time.Duration, exitchan chan error) error {
+// waitForSocket waits for the given file to exist
+func (m *Machine) waitForSocket(timeout time.Duration, exitchan chan error) error {
 	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	defer cancel()
-	exists := make(chan bool)
-	var err error
+
+	done := make(chan error)
+	ticker := time.NewTicker(10 * time.Millisecond)
+
 	go func() {
 		for {
-			_, err = os.Stat(path)
-			if err == nil {
-				exists <- true
+			select {
+			case <-ctx.Done():
+				done <- ctx.Err()
+				return
+			case err := <-exitchan:
+				done <- err
+				return
+			case <-ticker.C:
+				if _, err := os.Stat(m.cfg.SocketPath); err != nil {
+					continue
+				}
+
+				// Send test HTTP request to make sure socket is available
+				if _, err := m.client.GetMachineConfig(); err != nil {
+					continue
+				}
+
+				done <- nil
 				return
 			}
-			time.Sleep(10 * time.Millisecond)
 		}
 	}()
-	select {
-	case <-ctx.Done():
-		return ctx.Err()
-	case <-exists:
-		return nil
-	case <-exitchan:
-		return errors.New("Firecracker exited unexpectedly")
-	}
+
+	return <-done
 }

machine_mock_test.go

@@ -11,6 +11,8 @@
 // express or implied. See the License for the specific language governing
 // permissions and limitations under the License.
 
+//go:generate mockgen -source=machine.go -destination=machine_mock_test.go -package=firecracker
+
 package firecracker
 
 import (
@@ -21,6 +23,8 @@ import (
 	"path/filepath"
 	"testing"
 	"time"
+
+	"github.com/golang/mock/gomock"
 )
 
 const (
@@ -216,38 +220,63 @@ func testStopVMM(ctx context.Context, t *testing.T, m *Machine) {
 }
 
 func TestWaitForSocket(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	okClient := NewMockFirecracker(ctrl)
+	okClient.EXPECT().GetMachineConfig().AnyTimes().Return(nil, nil)
+
+	errClient := NewMockFirecracker(ctrl)
+	errClient.EXPECT().GetMachineConfig().AnyTimes().Return(nil, errors.New("http error"))
+
 	// testWaitForSocket has three conditions that need testing:
-	// 1. The expected file is created within the deadline
+	// 1. The expected file is created within the deadline and
+	//    socket HTTP request succeeded
 	// 2. The expected file is not created within the deadline
 	// 3. The process responsible for creating the file exits
-	//  (indicated by an error published to exitchan)
+	//    (indicated by an error published to exitchan)
 	filename := "./test-create-file"
 	errchan := make(chan error)
 
+	m := Machine{
+		cfg: Config{SocketPath: filename},
+	}
+
 	go func() {
 		time.Sleep(50 * time.Millisecond)
 		_, err := os.Create(filename)
 		if err != nil {
 			t.Fatalf("Unable to create test file %s: %s", filename, err)
 		}
 	}()
-	err := waitForSocket(filename, 500*time.Millisecond, errchan)
-	if err != nil {
+
+	// Socket file created, HTTP request succeeded
+	m.client = okClient
+	if err := m.waitForSocket(500*time.Millisecond, errchan); err != nil {
 		t.Errorf("waitForSocket returned unexpected error %s", err)
 	}
 
+	// Socket file exists, HTTP request failed
+	m.client = errClient
+	if err := m.waitForSocket(500*time.Millisecond, errchan); err != context.DeadlineExceeded {
+		t.Error("waitforSocket did not return an expected timeout error")
+	}
+
 	os.Remove(filename)
-	err = waitForSocket(filename, 100*time.Millisecond, errchan)
-	if err == nil {
+
+	// No socket file
+	if err := m.waitForSocket(100*time.Millisecond, errchan); err != context.DeadlineExceeded {
 		t.Error("waitforSocket did not return an expected timeout error")
 	}
 
+	chanErr := errors.New("this is an expected error")
 	go func() {
 		time.Sleep(50 * time.Millisecond)
-		errchan <- errors.New("This is an expected error")
+		errchan <- chanErr
 	}()
-	err = waitForSocket(filename, 100*time.Millisecond, errchan)
-	if err == nil {
+
+	// Unexpected process exit
+	if err := m.waitForSocket(100*time.Millisecond, errchan); err != chanErr {
 		t.Error("waitForSocket did not properly detect program exit")
 	}
 }