import com.sun.tools.attach.VirtualMachine;

import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

public class VThreadTls {
    public static native long getTls();
    public static native void setTls(long value);
    
    public static volatile boolean attached;
    public static volatile boolean error;

    public static void main(String[] args) throws Exception {
        try (ExecutorService executorService = Executors.newVirtualThreadPerTaskExecutor()) {
            for (int threadCount = 0; threadCount < 100; threadCount++) {
                executorService.execute(() -> {
                    try {
                        while (!attached) {
                            // keep mounted
                        }
                        long threadId = Thread.currentThread().threadId();
                        setTls(threadId);
                        long mountedValue = getTls();
                        if (mountedValue != threadId) {
                            System.out.println("wrong tls value while still mounted: " + threadId + ", " + mountedValue);
                            error = true;
                            return;
                        }

                        for (int repetion = 0; repetion < 1000; repetion++) {
                            Thread.sleep(1);
                            long tlsValue = getTls();
                            if (tlsValue != threadId) {
                                System.out.println("wrong tls value after yield: " + threadId + ", " + tlsValue);
                                error = true;
                                return;
                            }
                        }
                    } catch (InterruptedException e) {
                        throw new RuntimeException(e);
                    }
                });
            }
            if (args.length == 1 && args[0].equals("attach")) {
                VirtualMachine vm = VirtualMachine.attach(String.valueOf(ProcessHandle.current().pid()));
                vm.loadAgentLibrary("VThreadTls");
            }
            Thread.sleep(500);
            attached = true;
        }
        if (error) {
            System.exit(1);
        }
    }
}