// Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
package com.intellij.execution.testDiscovery;

import com.intellij.execution.process.ProcessAdapter;
import com.intellij.execution.process.ProcessEvent;
import com.intellij.execution.process.ProcessHandler;
import com.intellij.openapi.application.ApplicationManager;
import com.intellij.openapi.diagnostic.Logger;
import com.intellij.openapi.project.Project;
import com.intellij.rt.coverage.data.api.TestDiscoveryProtocolUtil;
import com.intellij.util.TimeoutUtil;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.jetbrains.annotations.TestOnly;

import java.io.IOException;
import java.io.InputStream;
import java.net.ServerSocket;
import java.net.Socket;

public class TestDiscoveryDataSocketListener {
  private static final Logger LOG = Logger.getInstance(TestDiscoveryDataSocketListener.class);

  @Nullable
  private final String myModuleName;
  private final byte myFrameworkId;
  private volatile boolean myClosed;
  private volatile boolean myFinished;
  private final ServerSocket myServer;
  private final int myPort;
  private final TestDiscoveryIndex myTestDiscoveryIndex;

  public TestDiscoveryDataSocketListener(@NotNull Project project,
                                         @Nullable String moduleName,
                                         byte frameworkPrefix) throws IOException {
    myTestDiscoveryIndex = TestDiscoveryIndex.getInstance(project);
    myModuleName = moduleName;
    myFrameworkId = frameworkPrefix;
    myServer = new ServerSocket(0);
    myPort = myServer.getLocalPort();

    ApplicationManager.getApplication().executeOnPooledThread(() -> {
      Socket socket;
      while (true) {
        if (myClosed) {
          return;
        }
        try {
          socket = myServer.accept();
          if (socket != null) {
            break;
          }
        }
        catch (IOException e) {
          LOG.error(e);
          return;
        }
        TimeoutUtil.sleep(10);
      }

      try {
        InputStream testDataStream = socket.getInputStream();
        IdeaTestDiscoveryProtocolReader protocolReader = new IdeaTestDiscoveryProtocolReader(myTestDiscoveryIndex, myModuleName,
                                                                                             myFrameworkId);
        TestDiscoveryProtocolUtil.readSequentially(testDataStream, protocolReader);
        myFinished = true;
      }
      catch (IOException e) {
        if (!myClosed) {
          LOG.error(e);
        }
      } finally {
        try {
          socket.close();
        }
        catch (IOException e) {
          LOG.error(e);
        }
      }
    });
  }

  public int getPort() {
    return myPort;
  }

  void attach(ProcessHandler handler) {
    handler.addProcessListener(new ProcessAdapter() {
      @Override
      public void processTerminated(@NotNull ProcessEvent event) {
        myClosed = true;
      }

      @Override
      public void processWillTerminate(@NotNull ProcessEvent event, boolean willBeDestroyed) {
        myClosed = true;
      }
    });
  }

  @TestOnly
  public void awaitTermination() {
    while (!myFinished) {
      TimeoutUtil.sleep(100);
    }
  }
}

@@ -1774,52 +1774,8 @@
 %3E %7B%0A
-      Socket socket;%0A      while (true) %7B%0A  
     
@@ -1792,34 +1792,32 @@
 osed) %7B%0A        
- 
- 
 return;%0A        
@@ -1810,30 +1810,54 @@
 turn;%0A      
-  
 %7D%0A  
+    Socket socket = null;%0A
       try %7B%0A
@@ -1848,34 +1848,32 @@
 ll;%0A      try %7B%0A
-  
         socket =
@@ -1904,208 +1904,85 @@
     
-  if (socket != null) %7B%0A            break;%0A          %7D%0A        %7D%0A        catch (IOException e) %7B%0A          LOG.error(e);%0A          return;%0A        %7D%0A        TimeoutUtil.sleep(10);%0A      %7D%0A%0A      try %7B
+socket.setReceiveBufferSize(1024 * 1024);%0A        socket.setTcpNoDelay(true);
 %0A   
@@ -2510,32 +2510,66 @@
 %7B%0A        try %7B%0A
+          if (socket != null) %7B%0A  
           socket
@@ -2574,24 +2574,36 @@
 et.close();%0A
+          %7D%0A
         %7D%0A  

test discovery: remove redundant loop in socket listener